check for input and output types for proofmode

cmd/cli/main.go

@@ -220,7 +220,7 @@ func main() {
 					if err != nil {
 						return fmt.Errorf("cannot parse args: %w", err)
 					}
-					program, hints, userArgs, err := runner.AssembleProgram(cairoProgram, userArgs, availableGas)
+					program, hints, userArgs, err := runner.AssembleProgram(cairoProgram, userArgs, availableGas, proofmode)
 					if err != nil {
 						return fmt.Errorf("cannot assemble program: %w", err)
 					}

integration_tests/.env

@@ -1,2 +1,2 @@
 # Set to run some specific file tests (ex. fib.cairo,alloc.cairo)
-INTEGRATION_TESTS_FILTERS=
+INTEGRATION_TESTS_FILTERS=proofmode__small.cairo

integration_tests/cairo_1_programs/with_input/proofmode__small.cairo

@@ -0,0 +1,22 @@
+fn arr_sum(mut vals: Array<felt252>) -> felt252 {
+    let mut sum: felt252 = 0;
+
+    loop {
+        match vals.pop_front() {
+            Option::Some(v) => {
+                sum += v;
+            },
+            Option::None(_) => {
+                break sum;
+            }
+        };
+    }
+}
+
+fn main(vals: Array<felt252>) -> Array<felt252> {
+    let sum1 = arr_sum(vals);
+    let mut res: Array<felt252> = ArrayTrait::new();
+    res.append(sum1);
+    res.append(sum1);
+    res
+}

integration_tests/cairo_vm_test.go

@@ -192,6 +192,7 @@ func TestCairoFiles(t *testing.T) {
 		"cairo_1_programs/with_input/branching.cairo":              "123",
 		"cairo_1_programs/with_input/dict_with_input__small.cairo": "[1 2 3 4]",
 		"cairo_1_programs/with_input/tensor__small.cairo":          "[1 4] [1 5]",
+		"cairo_1_programs/with_input/proofmode__small.cairo":       "[1 2 3 4 5]",
 	}
 
 	// filter is for debugging purposes
@@ -472,7 +473,7 @@ func runVm(path, layout string, zero bool, inputArgs string) (time.Duration, str
 	if !zero {
 		args = []string{
 			cliCommand,
-			// "--proofmode",
+			"--proofmode",
 			"--tracefile",
 			traceOutput,
 			"--memoryfile",

pkg/runner/runner.go

@@ -65,11 +65,17 @@ func NewRunner(program *Program, hints map[uint64][]hinter.Hinter, runnerMode Ru
 	}, nil
 }
 
-func AssembleProgram(cairoProgram *starknet.StarknetProgram, userArgs []starknet.CairoFuncArgs, availableGas uint64) (Program, map[uint64][]hinter.Hinter, []starknet.CairoFuncArgs, error) {
+func AssembleProgram(cairoProgram *starknet.StarknetProgram, userArgs []starknet.CairoFuncArgs, availableGas uint64, proofmode bool) (Program, map[uint64][]hinter.Hinter, []starknet.CairoFuncArgs, error) {
 	mainFunc, ok := cairoProgram.EntryPointsByFunction["main"]
 	if !ok {
 		return Program{}, nil, nil, fmt.Errorf("cannot find main function")
 	}
+	if proofmode {
+		err := CheckOnlyArrayFeltInputAndReturntValue(mainFunc)
+		if err != nil {
+			return Program{}, nil, nil, err
+		}
+	}
 	expectedArgsSize, actualArgsSize := 0, 0
 	for _, arg := range mainFunc.InputArgs {
 		expectedArgsSize += arg.Size
@@ -647,3 +653,19 @@ func GetFooterInstructions() []*fp.Element {
 	// and `pc` registers.
 	return []*fp.Element{new(fp.Element).SetUint64(2345108766317314046)}
 }
+
+func CheckOnlyArrayFeltInputAndReturntValue(mainFunc starknet.EntryPointByFunction) error {
+	if len(mainFunc.InputArgs) != 1 {
+		return fmt.Errorf("main function should have only one input argument")
+	}
+	if len(mainFunc.ReturnArg) != 1 {
+		return fmt.Errorf("main function should have only one return argument")
+	}
+	if mainFunc.InputArgs[0].Size != 2 || mainFunc.InputArgs[0].DebugName != "Array<felt252>" {
+		return fmt.Errorf("main function input argument should be Felt Array")
+	}
+	if mainFunc.ReturnArg[0].Size != 3 || mainFunc.ReturnArg[0].DebugName != "core::panics::PanicResult::<(core::array::Array::<core::felt252>)>" {
+		return fmt.Errorf("main function return argument should be Felt Array")
+	}
+	return nil
+}