hintrunner/zero: include the ap tracking data into addr calculations

This patch is aimed to include the ap tracking offsets into
the produced ApCellRef. The calculation result is saved
as an updated ApCellRef offset.

Given the `ApCellRef(-3)` and ap offset of hint `3` and ref `2`,
we'll get a new `ApCellRef(-4)` (-3 - (3 - 2) => -4).

Fixes #197

integration_tests/cairo_files/hintrefs.cairo

@@ -0,0 +1,64 @@
+// This test uses the artificially constructed TestAssignCode hint
+// to test different hint refs evaluation.
+//
+// Even if the hint's code is the same in every function below,
+// the referenced ids.a always has an address that requires
+// a different way of computation (see per-func comments).
+// They also usually have different ApTracking values associated with them.
+//
+// See https://github.com/NethermindEth/cairo-vm-go/issues/197
+
+// [cast(fp, felt*)]
+func simple_fp_ref() -> felt {
+    alloc_locals;
+    local a = 43;
+    %{ memory[ap] = ids.a %}
+    return [ap];
+}
+
+// [cast(ap + (-1), felt*)]
+func ap_with_offset() -> felt {
+    [ap] = 0, ap++;
+    [ap] = 10, ap++;
+    tempvar a = 32;
+    [ap] = 100, ap++;
+    [ap] = 200, ap++;
+    %{ memory[ap] = ids.a %}
+    return [ap];
+}
+
+// cast([fp + (-4)] + [fp + (-3)], felt)
+func fp_args_sum(arg1: felt, arg2: felt) -> felt {
+    let a = arg1 + arg2;
+    %{ memory[ap] = ids.a %}
+    return [ap];
+}
+
+// cast([ap + (-1)] + [fp + 1], felt)
+func ap_plus_fp_deref() -> felt {
+    alloc_locals;
+    local l1 = 11;
+    local l2 = 22; // [fp+1]
+    local l3 = 33;
+    tempvar t1 = 111;
+    tempvar t2 = 222;
+    tempvar t3 = 333; // [ap-1]
+    let a = [ap-1] + [fp+1];
+    %{ memory[ap] = ids.a %}
+    return [ap]; // 355
+}
+
+func main() {
+    alloc_locals;
+
+    local v1 = simple_fp_ref();
+    [ap] = v1, ap++;
+    local v2 = ap_with_offset();
+    [ap] = v2, ap++;
+    local v3 = fp_args_sum(4, 6);
+    [ap] = v3, ap++;
+    local v4 = ap_plus_fp_deref();
+    [ap] = v4, ap++;
+
+    ret;
+}

pkg/hintrunner/zero/hintcode.go

@@ -2,4 +2,9 @@ package zero
 
 const (
 	AllocSegmentCode string = "memory[ap] = segments.add()"
+
+	// This is a very simple Cairo0 hint that allows us to test
+	// the identifier resolution code.
+	// Depending on the context, ids.a may be a complex reference.
+	TestAssignCode string = "memory[ap] = ids.a"
 )

pkg/hintrunner/zero/zerohint.go

@@ -8,8 +8,23 @@ import (
 	"github.com/NethermindEth/cairo-vm-go/pkg/hintrunner/hinter"
 	sn "github.com/NethermindEth/cairo-vm-go/pkg/parsers/starknet"
 	zero "github.com/NethermindEth/cairo-vm-go/pkg/parsers/zero"
+	VM "github.com/NethermindEth/cairo-vm-go/pkg/vm"
 )
 
+// GenericZeroHinter wraps an adhoc Cairo0 inline (pythonic) hint implementation.
+type GenericZeroHinter struct {
+	Name string
+	Op   func(vm *VM.VirtualMachine, _ *hinter.HintRunnerContext) error
+}
+
+func (hint *GenericZeroHinter) String() string {
+	return hint.Name
+}
+
+func (hint *GenericZeroHinter) Execute(vm *VM.VirtualMachine, ctx *hinter.HintRunnerContext) error {
+	return hint.Op(vm, ctx)
+}
+
 func GetZeroHints(cairoZeroJson *zero.ZeroProgram) (map[uint64][]hinter.Hinter, error) {
 	hints := make(map[uint64][]hinter.Hinter)
 	for counter, rawHints := range cairoZeroJson.Hints {
@@ -40,6 +55,8 @@ func GetHintFromCode(program *zero.ZeroProgram, rawHint zero.Hint, hintPC uint64
 	switch rawHint.Code {
 	case AllocSegmentCode:
 		return CreateAllocSegmentHinter(cellRefParams, resOpParams)
+	case TestAssignCode:
+		return createTestAssignHinter(cellRefParams, resOpParams)
 	default:
 		return nil, fmt.Errorf("Not identified hint")
 	}
@@ -52,6 +69,52 @@ func CreateAllocSegmentHinter(cellRefParams []hinter.CellRefer, resOpParams []hi
 	return &core.AllocSegment{Dst: hinter.ApCellRef(0)}, nil
 }
 
+func createTestAssignHinter(cellRefParams []hinter.CellRefer, resOpParams []hinter.ResOperander) (hinter.Hinter, error) {
+	if len(resOpParams) < 1 {
+		return nil, fmt.Errorf("Expected at least 1 ResOperander")
+	}
+	if len(cellRefParams) != 0 {
+		return nil, fmt.Errorf("Expected 0 CellRefers (got %d)", len(cellRefParams))
+	}
+
+	// Given a Cairo0 code like this:
+	//
+	//	func fp_args_sum(arg1: felt, arg2: felt) -> felt {
+	//		let a = arg1 + arg2;
+	//		%{ memory[ap] = ids.a %}
+	//		return [ap];
+	//	}
+	//
+	// We get ids.a defined as a reference that refers to 2 function arguments.
+	// When we execute GetParameters(), it will return 3 ResOperander for the hint:
+	// one for the ids.a itself, and two for the args (as they're referenced from it).
+	//
+	//	"__main__.fp_args_sum.a"="cast([fp + (-4)] + [fp + (-3)], felt)"
+	//	"__main__.fp_args_sum.arg1"="[cast(fp + (-4), felt*)]"
+	//	"__main__.fp_args_sum.arg2"="[cast(fp + (-3), felt*)]"
+	//
+	// It's not entirely clear how to handle this situation yet.
+	// It looks like we usually need exactly 1 reference per literal identifier (like ids.a),
+	// but we may not need references that are not present in the hint's code directly.
+	//
+	// For now, I'll just take the first one (for tests purposes) and leave this issue for discussion.
+	// Right now I want to test the ApTracking address calculation, not hint's reference collection.
+	arg := resOpParams[0]
+
+	h := &GenericZeroHinter{
+		Name: "TestAssign",
+		Op: func(vm *VM.VirtualMachine, _ *hinter.HintRunnerContext) error {
+			apAddr := vm.Context.AddressAp()
+			v, err := arg.Resolve(vm)
+			if err != nil {
+				return err
+			}
+			return vm.Memory.WriteToAddress(&apAddr, &v)
+		},
+	}
+	return h, nil
+}
+
 func GetParameters(zeroProgram *zero.ZeroProgram, hint zero.Hint, hintPC uint64) ([]hinter.CellRefer, []hinter.ResOperander, error) {
 	var cellRefParams []hinter.CellRefer
 	var resOpParams []hinter.ResOperander
@@ -84,6 +147,7 @@ func GetParameters(zeroProgram *zero.ZeroProgram, hint zero.Hint, hintPC uint64)
 		if err != nil {
 			return nil, nil, err
 		}
+		param = applyApTracking(zeroProgram, hint, reference, param)
 		switch result := param.(type) {
 		case hinter.CellRefer:
 			cellRefParams = append(cellRefParams, result)
@@ -96,3 +160,36 @@ func GetParameters(zeroProgram *zero.ZeroProgram, hint zero.Hint, hintPC uint64)
 
 	return cellRefParams, resOpParams, nil
 }
+
+func applyApTracking(p *zero.ZeroProgram, h zero.Hint, ref zero.Reference, v any) any {
+	// We can't make an inplace modification because the v's underlying type is not a pointer type.
+	// Therefore, we need to return it from the function.
+	// This makes this function less elegant: it requires type asserts, etc.
+
+	switch v := v.(type) {
+	case hinter.ApCellRef:
+		if h.FlowTrackingData.ApTracking.Group != ref.ApTrackingData.Group {
+			return v // Group mismatched: nothing to adjust
+		}
+		newOffset := v - hinter.ApCellRef(h.FlowTrackingData.ApTracking.Offset-ref.ApTrackingData.Offset)
+		return hinter.ApCellRef(newOffset)
+
+	case hinter.Deref:
+		v.Deref = applyApTracking(p, h, ref, v.Deref).(hinter.CellRefer)
+		return v
+
+	case hinter.DoubleDeref:
+		v.Deref = applyApTracking(p, h, ref, v.Deref).(hinter.CellRefer)
+		return v
+
+	case hinter.BinaryOp:
+		v.Lhs = applyApTracking(p, h, ref, v.Lhs).(hinter.CellRefer)
+		v.Rhs = applyApTracking(p, h, ref, v.Rhs).(hinter.ResOperander)
+		return v
+
+	default:
+		// This case covers type that we don't need to visit.
+		// E.g. FpCellRef, Immediate.
+		return v
+	}
+}