Add RunUntilPc & GetReturnValues (#302)

* Implement GenArg

* Remove recursive processing

* Add unit tests

* Fix test values

* Start fn

* Add RunFromEntryPoint

* Add test for RunFromEntryPoint

* Add unit tests

* Add comments

pkg/runners/cairo_runner.go

@@ -70,11 +70,11 @@ func NewCairoRunner(program vm.Program, layoutName string, proofMode bool) (*Cai
 
 // Performs the initialization step, returns the end pointer (pc upon which execution should stop)
 func (r *CairoRunner) Initialize() (memory.Relocatable, error) {
-	err := r.initializeBuiltins()
+	err := r.InitializeBuiltins()
 	if err != nil {
 		return memory.Relocatable{}, errors.New(err.Error())
 	}
-	r.initializeSegments()
+	r.InitializeSegments()
 	end, err := r.initializeMainEntrypoint()
 	if err == nil {
 		err = r.initializeVM()
@@ -84,7 +84,7 @@ func (r *CairoRunner) Initialize() (memory.Relocatable, error) {
 
 // Initializes builtin runners in accordance to the specified layout and
 // the builtins present in the running program.
-func (r *CairoRunner) initializeBuiltins() error {
+func (r *CairoRunner) InitializeBuiltins() error {
 	var builtinRunners []builtins.BuiltinRunner
 	programBuiltins := map[string]struct{}{}
 	for _, builtin := range r.Program.Builtins {
@@ -113,7 +113,7 @@ func (r *CairoRunner) initializeBuiltins() error {
 }
 
 // Creates program, execution and builtin segments
-func (r *CairoRunner) initializeSegments() {
+func (r *CairoRunner) InitializeSegments() {
 	// Program Segment
 	r.ProgramBase = r.Vm.Segments.AddSegment()
 	// Execution Segment
@@ -144,9 +144,9 @@ func (r *CairoRunner) initializeState(entrypoint uint, stack *[]memory.MaybeRelo
 
 // Initializes memory, initial register values & returns the end pointer (final pc) to run from a given pc offset
 // (entrypoint)
-func (r *CairoRunner) initializeFunctionEntrypoint(entrypoint uint, stack *[]memory.MaybeRelocatable, return_fp memory.Relocatable) (memory.Relocatable, error) {
+func (r *CairoRunner) initializeFunctionEntrypoint(entrypoint uint, stack *[]memory.MaybeRelocatable, return_fp memory.MaybeRelocatable) (memory.Relocatable, error) {
 	end := r.Vm.Segments.AddSegment()
-	*stack = append(*stack, *memory.NewMaybeRelocatableRelocatable(return_fp), *memory.NewMaybeRelocatableRelocatable(end))
+	*stack = append(*stack, return_fp, *memory.NewMaybeRelocatableRelocatable(end))
 	r.initialFp = r.executionBase
 	r.initialFp.Offset += uint(len(*stack))
 	r.initialAp = r.initialFp
@@ -188,7 +188,7 @@ func (r *CairoRunner) initializeMainEntrypoint() (memory.Relocatable, error) {
 		return memory.NewRelocatable(r.ProgramBase.SegmentIndex, r.ProgramBase.Offset+r.Program.End), nil
 	}
 
-	return_fp := r.Vm.Segments.AddSegment()
+	return_fp := *memory.NewMaybeRelocatableRelocatable(r.Vm.Segments.AddSegment())
 	return r.initializeFunctionEntrypoint(r.mainOffset, &stack, return_fp)
 }
 
@@ -581,3 +581,40 @@ func (runner *CairoRunner) GetExecutionResources() (ExecutionResources, error) {
 		BuiltinsInstanceCounter: builtinInstaceCounter,
 	}, nil
 }
+
+// TODO: Add verifySecure once its implemented
+/*
+Runs a cairo program from a give entrypoint, indicated by its pc offset, with the given arguments.
+If `verifySecure` is set to true, [verifySecureRunner] will be called to run extra verifications.
+`programSegmentSize` is only used by the [verifySecureRunner] function and will be ignored if `verifySecure` is set to false.
+Each arg can be either MaybeRelocatable, []MaybeRelocatable or [][]MaybeRelocatable
+*/
+func (runner *CairoRunner) RunFromEntrypoint(entrypoint uint, args []any, hintProcessor vm.HintProcessor) error {
+	stack := make([]memory.MaybeRelocatable, 0)
+	for _, arg := range args {
+		val, err := runner.Vm.Segments.GenArg(arg)
+		if err != nil {
+			return err
+		}
+		stack = append(stack, val)
+	}
+	returnFp := *memory.NewMaybeRelocatableFelt(lambdaworks.FeltZero())
+	end, err := runner.initializeFunctionEntrypoint(entrypoint, &stack, returnFp)
+	if err != nil {
+		return err
+	}
+	err = runner.initializeVM()
+	if err != nil {
+		return err
+	}
+	err = runner.RunUntilPC(end, hintProcessor)
+	if err != nil {
+		return err
+	}
+	err = runner.EndRun(false, false, hintProcessor)
+	if err != nil {
+		return err
+	}
+	// TODO: verifySecureRunner
+	return nil
+}

pkg/runners/cairo_runner_test.go

@@ -726,3 +726,36 @@ func TestRunFibonacciGetExecutionResources(t *testing.T) {
 		t.Errorf("Wong ExecutionResources.\n Expected : %+v, got: %+v", expectedExecutionResources, executionResources)
 	}
 }
+
+// This test will run the `fib` function in the fibonacci.json program
+func TestRunFromEntryPointFibonacci(t *testing.T) {
+	compiledProgram, _ := parser.Parse("../../cairo_programs/fibonacci.json")
+	programJson := vm.DeserializeProgramJson(compiledProgram)
+
+	entrypoint := programJson.Identifiers["__main__.fib"].PC
+	args := []any{
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltOne()),
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltOne()),
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(10)),
+	}
+	runner, _ := runners.NewCairoRunner(programJson, "all_cairo", false)
+	hintProcessor := hints.CairoVmHintProcessor{}
+
+	runner.InitializeBuiltins()
+	runner.InitializeSegments()
+	err := runner.RunFromEntrypoint(uint(entrypoint), args, &hintProcessor)
+
+	if err != nil {
+		t.Errorf("Running fib entrypoint failed with error %s", err.Error())
+	}
+
+	// Check result
+	res, err := runner.Vm.GetReturnValues(1)
+	if err != nil {
+		t.Errorf("Failed to fetch return values from fib with error %s", err.Error())
+	}
+	if len(res) != 1 || !reflect.DeepEqual(res[0], *memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(144))) {
+		t.Errorf("Wrong value returned by fib entrypoint.\n Expected [144], got: %+v", res)
+	}
+
+}

pkg/vm/memory/segments.go

@@ -1,6 +1,8 @@
 package memory
 
 import (
+	"errors"
+
 	"github.com/lambdaclass/cairo-vm.go/pkg/lambdaworks"
 )
 
@@ -186,3 +188,39 @@ func (m *MemorySegmentManager) GetFeltRange(start Relocatable, size uint) ([]lam
 	}
 	return feltRange, nil
 }
+
+/*
+Converts a generic argument into a MaybeRelocatable
+If the argument is a slice, it loads it into memory in a new segment and returns its base
+Accepts MaybeRelocatable, []MaybeRelocatable, [][]MaybeRelocatable
+*/
+func (m *MemorySegmentManager) GenArg(arg any) (MaybeRelocatable, error) {
+	// Attempt to cast to MaybeRelocatable
+	a, ok := arg.(MaybeRelocatable)
+	if ok {
+		return a, nil
+	}
+	// Attempt to cast to []MaybeRelocatable
+	data, ok := arg.([]MaybeRelocatable)
+	if ok {
+		base := m.AddSegment()
+		_, err := m.LoadData(base, &data)
+		return *NewMaybeRelocatableRelocatable(base), err
+	}
+	// Attempt to cast to [][]MaybeRelocatable
+	datas, ok := arg.([][]MaybeRelocatable)
+	if ok {
+		args := make([]MaybeRelocatable, 0)
+		for _, data = range datas {
+			dataBase, err := m.GenArg(data)
+			if err != nil {
+				return *NewMaybeRelocatableFelt(lambdaworks.FeltZero()), err
+			}
+			args = append(args, dataBase)
+		}
+		base := m.AddSegment()
+		_, err := m.LoadData(base, &args)
+		return *NewMaybeRelocatableRelocatable(base), err
+	}
+	return *NewMaybeRelocatableFelt(lambdaworks.FeltZero()), errors.New("GenArg: found argument of invalid type.")
+}

pkg/vm/memory/segments_test.go

@@ -310,3 +310,52 @@ func TestGetFeltRangeRelocatable(t *testing.T) {
 		t.Errorf("GetFeltRange should have failed")
 	}
 }
+
+func TestGenArgMaybeRelocatable(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	arg := any(*memory.NewMaybeRelocatableFelt(lambdaworks.FeltZero()))
+	expectedArg := *memory.NewMaybeRelocatableFelt(lambdaworks.FeltZero())
+	genedArg, err := segments.GenArg(arg)
+	if err != nil || !reflect.DeepEqual(expectedArg, genedArg) {
+		t.Error("GenArg failed or returned wrong value")
+	}
+}
+
+func TestGenArgSliceMaybeRelocatable(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	arg := any([]memory.MaybeRelocatable{*memory.NewMaybeRelocatableFelt(lambdaworks.FeltZero())})
+
+	expectedBase := memory.NewRelocatable(0, 0)
+	expectedArg := *memory.NewMaybeRelocatableRelocatable(expectedBase)
+	genedArg, err := segments.GenArg(arg)
+	if err != nil || !reflect.DeepEqual(expectedArg, genedArg) {
+		t.Error("GenArg failed or returned wrong value")
+	}
+	val, err := segments.Memory.GetFelt(expectedBase)
+	if err != nil || !val.IsZero() {
+		t.Error("GenArg inserted wrong value into memory")
+	}
+}
+
+func TestGenArgSliceSliceMaybeRelocatable(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	arg := any([][]memory.MaybeRelocatable{{*memory.NewMaybeRelocatableFelt(lambdaworks.FeltZero())}})
+
+	expectedBaseA := memory.NewRelocatable(1, 0)
+	expectedBaseB := memory.NewRelocatable(0, 0)
+	expectedArg := *memory.NewMaybeRelocatableRelocatable(expectedBaseA)
+	genedArg, err := segments.GenArg(arg)
+
+	if err != nil || !reflect.DeepEqual(expectedArg, genedArg) {
+		t.Error("GenArg failed or returned wrong value")
+	}
+	valA, err := segments.Memory.GetRelocatable(expectedBaseA)
+	if err != nil || valA != expectedBaseB {
+		t.Error("GenArg inserted wrong value into memory")
+	}
+
+	valB, err := segments.Memory.GetFelt(expectedBaseB)
+	if err != nil || !valB.IsZero() {
+		t.Error("GenArg inserted wrong value into memory")
+	}
+}

pkg/vm/vm_core.go

@@ -642,3 +642,12 @@ func (vm *VirtualMachine) GetRangeCheckBound() (lambdaworks.Felt, error) {
 
 	return rcBuiltin.Bound(), nil
 }
+
+// Gets `nRet` return values from memory
+func (vm *VirtualMachine) GetReturnValues(nRet uint) ([]memory.MaybeRelocatable, error) {
+	ptr, err := vm.RunContext.Ap.SubUint(nRet)
+	if err != nil {
+		return nil, err
+	}
+	return vm.Segments.Memory.GetRange(ptr, nRet)
+}

pkg/vm/vm_test.go

@@ -1088,3 +1088,31 @@ func TestGetFooBuiltinReturnsNilAndError(t *testing.T) {
 		t.Error("Obtained a non existant builtin, or didn't raise an error")
 	}
 }
+
+func TestReadReturnValuesOk(t *testing.T) {
+	vm := vm.NewVirtualMachine()
+	vm.Segments.AddSegment()
+	// Load data at ap and advance ap
+	data := []memory.MaybeRelocatable{
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(1)),
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(2)),
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(3)),
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(4)),
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(5)),
+	}
+	vm.RunContext.Ap, _ = vm.Segments.LoadData(vm.RunContext.Ap, &data)
+	// Fetch 3 return values
+	expectedReturnValues := []memory.MaybeRelocatable{
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(3)),
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(4)),
+		*memory.NewMaybeRelocatableFelt(lambdaworks.FeltFromUint(5)),
+	}
+	returnValues, err := vm.GetReturnValues(3)
+	if err != nil {
+		t.Errorf("GetReturnValues failed with error: %s", err.Error())
+	}
+
+	if !reflect.DeepEqual(expectedReturnValues, returnValues) {
+		t.Errorf("Wrong return values.\n Expected: %+v, got: %+v", expectedReturnValues, returnValues)
+	}
+}