Add functionalities for main function inputs

cmd/cli/main.go

@@ -12,7 +12,6 @@ import (
 	"github.com/NethermindEth/cairo-vm-go/pkg/parsers/starknet"
 	zero "github.com/NethermindEth/cairo-vm-go/pkg/parsers/zero"
 	"github.com/NethermindEth/cairo-vm-go/pkg/runner"
-	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 	"github.com/urfave/cli/v2"
 )
 
@@ -217,26 +216,18 @@ func main() {
 					if err != nil {
 						return fmt.Errorf("cannot load program: %w", err)
 					}
-					program, hints, err := runner.AssembleProgram(cairoProgram)
+					userArgs, err := starknet.ParseCairoProgramArgs(args)
+					if err != nil {
+						return fmt.Errorf("cannot parse args: %w", err)
+					}
+					program, hints, userArgs, err := runner.AssembleProgram(cairoProgram, userArgs, availableGas)
 					if err != nil {
 						return fmt.Errorf("cannot assemble program: %w", err)
 					}
 					runnerMode := runner.ExecutionModeCairo
 					if proofmode {
 						runnerMode = runner.ProofModeCairo
 					}
-					userArgs, err := starknet.ParseCairoProgramArgs(args)
-					if err != nil {
-						return fmt.Errorf("cannot parse args: %w", err)
-					}
-					if availableGas > 0 {
-						// The first argument is the available gas
-						availableGasArg := starknet.CairoFuncArgs{
-							Single: new(fp.Element).SetUint64(availableGas),
-							Array:  nil,
-						}
-						userArgs = append([]starknet.CairoFuncArgs{availableGasArg}, userArgs...)
-					}
 					return runVM(program, proofmode, maxsteps, entrypointOffset, collectTrace, traceLocation, buildMemory, memoryLocation, layoutName, airPublicInputLocation, airPrivateInputLocation, hints, runnerMode, userArgs)
 				},
 			},

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=poseidon_pedersen__starknet.cairo

integration_tests/cairo_vm_test.go

@@ -66,7 +66,7 @@ func (f *Filter) filtered(testFile string) bool {
 	return false
 }
 
-func runAndTestFile(t *testing.T, path string, name string, benchmarkMap map[string][3]int, benchmark bool, errorExpected bool, zero bool) {
+func runAndTestFile(t *testing.T, path string, name string, benchmarkMap map[string][3]int, benchmark bool, errorExpected bool, zero bool, inputArgs string) {
 	t.Logf("testing: %s\n", path)
 	compiledOutput, err := compileCairoCode(path, zero)
 	if err != nil {
@@ -75,7 +75,7 @@ func runAndTestFile(t *testing.T, path string, name string, benchmarkMap map[str
 	}
 	layout := getLayoutFromFileName(path)
 
-	elapsedGo, traceFile, memoryFile, _, err := runVm(compiledOutput, layout, zero)
+	elapsedGo, traceFile, memoryFile, _, err := runVm(compiledOutput, layout, zero, inputArgs)
 	if errorExpected {
 		assert.Error(t, err, path)
 		writeToFile(path)
@@ -92,7 +92,7 @@ func runAndTestFile(t *testing.T, path string, name string, benchmarkMap map[str
 	if zero {
 		rustVmFilePath = compiledOutput
 	}
-	elapsedRs, rsTraceFile, rsMemoryFile, err := runRustVm(rustVmFilePath, layout, zero)
+	elapsedRs, rsTraceFile, rsMemoryFile, err := runRustVm(rustVmFilePath, layout, zero, inputArgs)
 	if errorExpected {
 		// we let the code go on so that we can check if the go vm also raises an error
 		assert.Error(t, err, path)
@@ -171,15 +171,27 @@ func TestCairoFiles(t *testing.T) {
 	if err != nil {
 		panic(fmt.Errorf("failed to open file: %w", err))
 	}
+
 	file.Close()
-	roots := []struct {
+	type TestCase struct {
 		path string
 		zero bool
-	}{
-		{"./cairo_zero_hint_tests/", true},
-		{"./cairo_zero_file_tests/", true},
-		{"./builtin_tests/", true},
-		// {"./cairo_1_programs/", false},
+	}
+	roots := []TestCase{
+		// {"./cairo_zero_hint_tests/", true},
+		// {"./cairo_zero_file_tests/", true},
+		// {"./builtin_tests/", true},
+		{"./cairo_1_programs/", false},
+		{"./cairo_1_programs/dict_non_squashed", false},
+		{"./cairo_1_programs/with_input", false},
+	}
+
+	inputArgsMap := map[string]string{
+		"cairo_1_programs/with_input/array_input_sum__small.cairo": "2 [111 222 333] 1 [444 555 666 777]",
+		"cairo_1_programs/with_input/array_length__small.cairo":    "[1 2 3 4 5 6] [7 8 9 10]",
+		"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]",
 	}
 
 	// filter is for debugging purposes
@@ -210,22 +222,19 @@ func TestCairoFiles(t *testing.T) {
 			if !filter.filtered(name) {
 				continue
 			}
-
+			inputArgs := inputArgsMap[path]
 			// we run tests concurrently if we don't need benchmarks
 			if !*zerobench {
 				sem <- struct{}{} // acquire a semaphore slot
 				wg.Add(1)
 
-				go func(path, name string, root struct {
-					path string
-					zero bool
-				}) {
+				go func(path, name string, root TestCase, inputArgs string) {
 					defer wg.Done()
 					defer func() { <-sem }() // release the semaphore slot when done
-					runAndTestFile(t, path, name, benchmarkMap, *zerobench, errorExpected, root.zero)
-				}(path, name, root)
+					runAndTestFile(t, path, name, benchmarkMap, *zerobench, errorExpected, root.zero, inputArgs)
+				}(path, name, root, inputArgs)
 			} else {
-				runAndTestFile(t, path, name, benchmarkMap, *zerobench, errorExpected, root.zero)
+				runAndTestFile(t, path, name, benchmarkMap, *zerobench, errorExpected, root.zero, inputArgs)
 			}
 		}
 	}
@@ -399,7 +408,7 @@ func runPythonVm(path, layout string) (time.Duration, string, string, error) {
 
 // given a path to a compiled cairo zero file, execute it using the
 // rust vm and return the trace and memory files location
-func runRustVm(path, layout string, zero bool) (time.Duration, string, string, error) {
+func runRustVm(path, layout string, zero bool, inputArgs string) (time.Duration, string, string, error) {
 	traceOutput := swapExtenstion(path, rsTraceSuffix)
 	memoryOutput := swapExtenstion(path, rsMemorySuffix)
 
@@ -411,6 +420,8 @@ func runRustVm(path, layout string, zero bool) (time.Duration, string, string, e
 		memoryOutput,
 		"--layout",
 		layout,
+		"--args",
+		inputArgs,
 	}
 
 	if zero {
@@ -440,7 +451,7 @@ func runRustVm(path, layout string, zero bool) (time.Duration, string, string, e
 
 // given a path to a compiled cairo zero file, execute
 // it using our vm
-func runVm(path, layout string, zero bool) (time.Duration, string, string, string, error) {
+func runVm(path, layout string, zero bool, inputArgs string) (time.Duration, string, string, string, error) {
 	traceOutput := swapExtenstion(path, traceSuffix)
 	memoryOutput := swapExtenstion(path, memorySuffix)
 
@@ -470,6 +481,8 @@ func runVm(path, layout string, zero bool) (time.Duration, string, string, strin
 			layout,
 			"--available_gas",
 			"9999999",
+			"--args",
+			inputArgs,
 		}
 	}
 	args = append(args, path)

pkg/hintrunner/core/hint.go

@@ -1143,7 +1143,8 @@ func (hint *Felt252DictEntryInit) Execute(vm *VM.VirtualMachine, ctx *hinter.Hin
 
 	prevValue, err := ctx.DictionaryManager.At(dictPtr, key)
 	if err != nil {
-		return fmt.Errorf("get dictionary entry: %w", err)
+		mv := mem.MemoryValueFromFieldElement(&utils.FeltZero)
+		prevValue = &mv
 	}
 	if prevValue == nil {
 		mv := mem.EmptyMemoryValueAsFelt()
@@ -1237,7 +1238,8 @@ func (hint *InitSquashData) Execute(vm *VM.VirtualMachine, ctx *hinter.HintRunne
 	// todo(rodro): Don't know if it could be called multiple times, or
 	err := hinter.InitializeSquashedDictionaryManager(ctx)
 	if err != nil {
-		return err
+		ctx.SquashedDictionaryManager = hinter.SquashedDictionaryManager{}
+		_ = hinter.InitializeSquashedDictionaryManager(ctx)
 	}
 
 	dictAccessPtr, err := hinter.ResolveAsAddress(vm, hint.DictAccesses)
@@ -1272,7 +1274,7 @@ func (hint *InitSquashData) Execute(vm *VM.VirtualMachine, ctx *hinter.HintRunne
 
 	// sort the keys in descending order
 	sort.Slice(ctx.SquashedDictionaryManager.Keys, func(i, j int) bool {
-		return ctx.SquashedDictionaryManager.Keys[i].Cmp(&ctx.SquashedDictionaryManager.Keys[j]) < 0
+		return ctx.SquashedDictionaryManager.Keys[i].Cmp(&ctx.SquashedDictionaryManager.Keys[j]) > 0
 	})
 
 	// if the first key is bigger than 2^128, signal it
@@ -1401,11 +1403,15 @@ func (hint *ShouldContinueSquashLoop) Execute(vm *VM.VirtualMachine, ctx *hinter
 	}
 
 	var shouldContinueLoop f.Element
-	if lastIndices, err := ctx.SquashedDictionaryManager.LastIndices(); err == nil && len(lastIndices) <= 1 {
-		shouldContinueLoop.SetOne()
-	} else if err != nil {
+	lastIndices, err := ctx.SquashedDictionaryManager.LastIndices()
+	if err != nil {
 		return fmt.Errorf("get last indices: %w", err)
 	}
+	if len(lastIndices) > 1 {
+		shouldContinueLoop.SetOne()
+	} else {
+		shouldContinueLoop.SetZero()
+	}
 
 	mv := mem.MemoryValueFromFieldElement(&shouldContinueLoop)
 	return vm.Memory.WriteToAddress(&shouldContinuePtr, &mv)
@@ -1425,11 +1431,15 @@ func (hint *GetNextDictKey) Execute(vm *VM.VirtualMachine, ctx *hinter.HintRunne
 		return fmt.Errorf("get next key address: %w", err)
 	}
 
-	nextKey, err := ctx.SquashedDictionaryManager.PopKey()
+	_, err = ctx.SquashedDictionaryManager.PopKey()
 	if err != nil {
 		return fmt.Errorf("pop key: %w", err)
 	}
 
+	nextKey, err := ctx.SquashedDictionaryManager.LastKey()
+	if err != nil {
+		return fmt.Errorf("get last key: %w", err)
+	}
 	mv := mem.MemoryValueFromFieldElement(&nextKey)
 	return vm.Memory.WriteToAddress(&nextKeyAddr, &mv)
 }
@@ -1928,29 +1938,47 @@ func (hint *ExternalWriteArgsToMemory) String() string {
 }
 
 func (hint *ExternalWriteArgsToMemory) Execute(vm *VM.VirtualMachine, ctx *hinter.HintRunnerContext) error {
-	userArgsVar, err := ctx.ScopeManager.GetVariableValue("userArgs")
+	userArgs, err := hinter.GetVariableAs[[]starknet.CairoFuncArgs](&ctx.ScopeManager, "userArgs")
 	if err != nil {
 		return fmt.Errorf("get user args: %v", err)
 	}
-	userArgs, ok := userArgsVar.([]starknet.CairoFuncArgs)
-	if !ok {
-		return fmt.Errorf("expected user args to be a list of CairoFuncArgs")
+	apOffset, err := hinter.GetVariableAs[uint64](&ctx.ScopeManager, "apOffset")
+	if err != nil {
+		return fmt.Errorf("get ap offset: %v", err)
 	}
+	fmt.Println("apOffset", apOffset, "ap.offset", vm.Context.Ap)
+	apOffset += vm.Context.Ap
 	for _, arg := range userArgs {
 		if arg.Single != nil {
 			mv := mem.MemoryValueFromFieldElement(arg.Single)
-			err := vm.Memory.Write(1, vm.Context.Ap, &mv)
+			err := vm.Memory.Write(1, apOffset, &mv)
 			if err != nil {
 				return fmt.Errorf("write single arg: %v", err)
 			}
+			apOffset++
 		} else if arg.Array != nil {
 			arrayBase := vm.Memory.AllocateEmptySegment()
 			mv := mem.MemoryValueFromMemoryAddress(&arrayBase)
-			err := vm.Memory.Write(1, vm.Context.Ap, &mv)
+			err := vm.Memory.Write(1, apOffset, &mv)
 			if err != nil {
 				return fmt.Errorf("write array base: %v", err)
 			}
-			// TODO: Implement array writing
+			apOffset++
+			arrayEnd := arrayBase
+			for _, val := range arg.Array {
+				arrayEnd.Offset += 1
+				mv := mem.MemoryValueFromFieldElement(&val)
+				err := vm.Memory.Write(arrayEnd.SegmentIndex, arrayEnd.Offset, &mv)
+				if err != nil {
+					return fmt.Errorf("write array element: %v", err)
+				}
+			}
+			mv = mem.MemoryValueFromMemoryAddress(&arrayEnd)
+			err = vm.Memory.Write(1, apOffset, &mv)
+			if err != nil {
+				return fmt.Errorf("write array end: %v", err)
+			}
+			apOffset++
 		}
 	}
 	return nil

pkg/hintrunner/hintrunner.go

@@ -15,10 +15,13 @@ type HintRunner struct {
 	hints map[uint64][]h.Hinter
 }
 
-func NewHintRunner(hints map[uint64][]h.Hinter, userArgs []starknet.CairoFuncArgs) HintRunner {
+func NewHintRunner(hints map[uint64][]h.Hinter, userArgs []starknet.CairoFuncArgs, writeApOffset uint64) HintRunner {
 	context := *h.InitializeDefaultContext()
 	if userArgs != nil {
-		err := context.ScopeManager.AssignVariable("userArgs", userArgs)
+		err := context.ScopeManager.AssignVariables(map[string]any{
+			"userArgs": userArgs,
+			"apOffset": writeApOffset,
+		})
 		// Error handling: this condition should never be true, since the context was initialized above
 		if err != nil {
 			panic(fmt.Errorf("assign userArgs: %v", err))

pkg/hintrunner/hintrunner_test.go

@@ -20,7 +20,7 @@ func TestExistingHint(t *testing.T) {
 
 	hr := NewHintRunner(map[uint64][]hinter.Hinter{
 		10: {&allocHint},
-	}, nil)
+	}, nil, 0)
 
 	vm.Context.Pc = memory.MemoryAddress{
 		SegmentIndex: 0,
@@ -44,7 +44,7 @@ func TestNoHint(t *testing.T) {
 
 	hr := NewHintRunner(map[uint64][]hinter.Hinter{
 		10: {&allocHint},
-	}, nil)
+	}, nil, 0)
 
 	vm.Context.Pc = memory.MemoryAddress{
 		SegmentIndex: 0,

pkg/runner/gas.go

@@ -51,7 +51,7 @@ func gasInitialization(memory *mem.Memory) error {
 		return err
 	}
 
-	preCostTokenTypes := []TokenGasCost{PedersenToken, BitwiseToken, EcOpToken, PoseidonToken, AddModToken, MulModToken}
+	preCostTokenTypes := []TokenGasCost{PedersenToken, PoseidonToken, BitwiseToken, EcOpToken, AddModToken, MulModToken}
 
 	for _, token := range preCostTokenTypes {
 		cost, err := getTokenGasCost(token)