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

…198)

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/hint_reference_resolver.go

@@ -0,0 +1,56 @@
+package zero
+
+import (
+	"github.com/NethermindEth/cairo-vm-go/pkg/hintrunner/hinter"
+)
+
+type hintReferenceResolver struct {
+	refs []hintReference
+
+	numCellRefer    int
+	numResOperander int
+}
+
+type hintReference struct {
+	name string
+
+	cell      hinter.CellRefer
+	operander hinter.ResOperander
+}
+
+func (m *hintReferenceResolver) NumCellRefers() int { return m.numCellRefer }
+
+func (m *hintReferenceResolver) NumResOperanders() int { return m.numResOperander }
+
+func (m *hintReferenceResolver) AddCellRefer(name string, v hinter.CellRefer) {
+	m.refs = append(m.refs, hintReference{
+		name: name,
+		cell: v,
+	})
+	m.numCellRefer++
+}
+
+func (m *hintReferenceResolver) AddResOperander(name string, v hinter.ResOperander) {
+	m.refs = append(m.refs, hintReference{
+		name:      name,
+		operander: v,
+	})
+	m.numResOperander++
+}
+
+func (m *hintReferenceResolver) GetResOperander(name string) hinter.ResOperander {
+	ref := m.find(name)
+	if ref != nil {
+		return ref.operander
+	}
+	return nil
+}
+
+func (m *hintReferenceResolver) find(name string) *hintReference {
+	for i, ref := range m.refs {
+		if ref.name == name {
+			return &m.refs[i]
+		}
+	}
+	return nil
+}

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

@@ -3,13 +3,29 @@ package zero
 import (
 	"fmt"
 	"strconv"
+	"strings"
 
 	"github.com/NethermindEth/cairo-vm-go/pkg/hintrunner/core"
 	"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 {
@@ -32,37 +48,69 @@ func GetZeroHints(cairoZeroJson *zero.ZeroProgram) (map[uint64][]hinter.Hinter,
 }
 
 func GetHintFromCode(program *zero.ZeroProgram, rawHint zero.Hint, hintPC uint64) (hinter.Hinter, error) {
-	cellRefParams, resOpParams, err := GetParameters(program, rawHint, hintPC)
+	resolver, err := getParameters(program, rawHint, hintPC)
 	if err != nil {
 		return nil, err
 	}
 
 	switch rawHint.Code {
 	case AllocSegmentCode:
-		return CreateAllocSegmentHinter(cellRefParams, resOpParams)
+		return CreateAllocSegmentHinter(resolver)
+	case TestAssignCode:
+		return createTestAssignHinter(resolver)
 	default:
 		return nil, fmt.Errorf("Not identified hint")
 	}
 }
 
-func CreateAllocSegmentHinter(cellRefParams []hinter.CellRefer, resOpParams []hinter.ResOperander) (hinter.Hinter, error) {
-	if len(cellRefParams)+len(resOpParams) != 0 {
+func CreateAllocSegmentHinter(resolver hintReferenceResolver) (hinter.Hinter, error) {
+	if resolver.NumResOperanders()+resolver.NumCellRefers() != 0 {
 		return nil, fmt.Errorf("Expected no arguments for %s hint", sn.AllocSegmentName)
 	}
 	return &core.AllocSegment{Dst: hinter.ApCellRef(0)}, nil
 }
 
-func GetParameters(zeroProgram *zero.ZeroProgram, hint zero.Hint, hintPC uint64) ([]hinter.CellRefer, []hinter.ResOperander, error) {
-	var cellRefParams []hinter.CellRefer
-	var resOpParams []hinter.ResOperander
+func createTestAssignHinter(resolver hintReferenceResolver) (hinter.Hinter, error) {
+	if resolver.NumResOperanders() < 1 {
+		return nil, fmt.Errorf("Expected at least 1 ResOperander")
+	}
+
+	arg := resolver.GetResOperander("a")
+
+	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
+}
+
+// shortSymbolName turns a full symbol name like "a.b.c" into just "c".
+func shortSymbolName(name string) string {
+	i := strings.LastIndexByte(name, '.')
+	if i != -1 {
+		return name[i+1:]
+	}
+	return name
+}
+
+func getParameters(zeroProgram *zero.ZeroProgram, hint zero.Hint, hintPC uint64) (hintReferenceResolver, error) {
+	var resolver hintReferenceResolver
+
 	for referenceName := range hint.FlowTrackingData.ReferenceIds {
 		rawIdentifier, ok := zeroProgram.Identifiers[referenceName]
 		if !ok {
-			return nil, nil, fmt.Errorf("missing identifier %s", referenceName)
+			return resolver, fmt.Errorf("missing identifier %s", referenceName)
 		}
 
 		if len(rawIdentifier.References) == 0 {
-			return nil, nil, fmt.Errorf("identifier %s should have at least one reference", referenceName)
+			return resolver, fmt.Errorf("identifier %s should have at least one reference", referenceName)
 		}
 		references := rawIdentifier.References
 
@@ -77,22 +125,56 @@ func GetParameters(zeroProgram *zero.ZeroProgram, hint zero.Hint, hintPC uint64)
 			}
 		}
 		if !ok {
-			return nil, nil, fmt.Errorf("identifier %s should have a reference with pc smaller or equal than %d", referenceName, hintPC)
+			return resolver, fmt.Errorf("identifier %s should have a reference with pc smaller or equal than %d", referenceName, hintPC)
 		}
 
 		param, err := ParseIdentifier(reference.Value)
 		if err != nil {
-			return nil, nil, err
+			return resolver, err
 		}
+		param = applyApTracking(zeroProgram, hint, reference, param)
 		switch result := param.(type) {
 		case hinter.CellRefer:
-			cellRefParams = append(cellRefParams, result)
+			resolver.AddCellRefer(shortSymbolName(referenceName), result)
 		case hinter.ResOperander:
-			resOpParams = append(resOpParams, result)
+			resolver.AddResOperander(shortSymbolName(referenceName), result)
 		default:
-			return nil, nil, fmt.Errorf("unexpected type for identifier value %s", reference.Value)
+			return resolver, fmt.Errorf("unexpected type for identifier value %s", reference.Value)
 		}
 	}
 
-	return cellRefParams, resOpParams, nil
+	return resolver, 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
+	}
 }