feat: SquashDictInnerNextKey hint (#400)

* SquashDictInnerNextKey hint * lint * address comments * address comments * fix * rectify pointer error
NethermindEth · May 10, 2024 · 58ca6fa · 58ca6fa
1 parent 1033f3e
commit 58ca6fa
Show file tree

Hide file tree

Showing 5 changed files with 120 additions and 3 deletions.
diff --git a/pkg/hintrunner/zero/hintcode.go b/pkg/hintrunner/zero/hintcode.go
@@ -88,6 +88,7 @@ const (
 
 	// ------ Dictionaries hints related code ------
 	squashDictInnerAssertLenKeys string = "assert len(keys) == 0"
+	squashDictInnerNextKey       string = "assert len(keys) > 0, 'No keys left but remaining_accesses > 0.'\nids.next_key = key = keys.pop()"
 
 	// ------ Other hints related code ------
 	allocSegmentCode     string = "memory[ap] = segments.add()"

diff --git a/pkg/hintrunner/zero/zerohint.go b/pkg/hintrunner/zero/zerohint.go
@@ -155,6 +155,8 @@ func GetHintFromCode(program *zero.ZeroProgram, rawHint zero.Hint, hintPC uint64
 		// Dictionaries hints
 	case squashDictInnerAssertLenKeys:
 		return createSquashDictInnerAssertLenKeysHinter()
+	case squashDictInnerNextKey:
+		return createSquashDictInnerNextKeyHinter(resolver)
 		// Other hints
 	case allocSegmentCode:
 		return createAllocSegmentHinter()

diff --git a/pkg/hintrunner/zero/zerohint_dictionaries.go b/pkg/hintrunner/zero/zerohint_dictionaries.go
@@ -4,7 +4,9 @@ import (
 	"fmt"
 
 	"github.com/NethermindEth/cairo-vm-go/pkg/hintrunner/hinter"
+	"github.com/NethermindEth/cairo-vm-go/pkg/utils"
 	VM "github.com/NethermindEth/cairo-vm-go/pkg/vm"
+	"github.com/NethermindEth/cairo-vm-go/pkg/vm/memory"
 	f "github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 )
 
@@ -33,3 +35,61 @@ func newSquashDictInnerAssertLenKeysHint() hinter.Hinter {
 func createSquashDictInnerAssertLenKeysHinter() (hinter.Hinter, error) {
 	return newSquashDictInnerAssertLenKeysHint(), nil
 }
+
+// SquashDictInnerNextKey hint retrieves the next key for processing during
+// dictionary squashing after checking that the array of keys is not empty
+//
+// `newSquashDictInnerNextKeyHint` takes 1 operander as argument
+//   - `next_key` variable will be assigned to the next key in `keys`
+func newSquashDictInnerNextKeyHint(nextKey hinter.ResOperander) hinter.Hinter {
+	return &GenericZeroHinter{
+		Name: "SquashDictInnerNextKey",
+		Op: func(vm *VM.VirtualMachine, ctx *hinter.HintRunnerContext) error {
+			//> assert len(keys) > 0, 'No keys left but remaining_accesses > 0.'
+			//> ids.next_key = key = keys.pop()
+
+			keys_, err := ctx.ScopeManager.GetVariableValue("keys")
+			if err != nil {
+				return err
+			}
+
+			keys := keys_.([]f.Element)
+			if len(keys) == 0 {
+				return fmt.Errorf("no keys left but remaining_accesses > 0")
+			}
+
+			newKey, err := utils.Pop(&keys)
+			if err != nil {
+				return err
+			}
+
+			err = ctx.ScopeManager.AssignVariable("keys", keys)
+			if err != nil {
+				return err
+			}
+
+			err = ctx.ScopeManager.AssignVariable("key", newKey)
+			if err != nil {
+				return err
+			}
+
+			newKeyMemoryValue := memory.MemoryValueFromFieldElement(&newKey)
+
+			addrNextKey, err := nextKey.GetAddress(vm)
+			if err != nil {
+				return err
+			}
+
+			return vm.Memory.WriteToAddress(&addrNextKey, &newKeyMemoryValue)
+		},
+	}
+}
+
+func createSquashDictInnerNextKeyHinter(resolver hintReferenceResolver) (hinter.Hinter, error) {
+	nextKey, err := resolver.GetResOperander("next_key")
+	if err != nil {
+		return nil, err
+	}
+
+	return newSquashDictInnerNextKeyHint(nextKey), nil
+}
diff --git a/pkg/hintrunner/zero/zerohint_dictionaries_test.go b/pkg/hintrunner/zero/zerohint_dictionaries_test.go
@@ -37,5 +37,58 @@ func TestZeroHintDictionaries(t *testing.T) {
 				errCheck: errorTextContains("assertion `len(keys) == 0` failed"),
 			},
 		},
+		"SquashDictInnerNextKey": {
+			{
+				operanders: []*hintOperander{
+					{Name: "next_key", Kind: uninitialized},
+				},
+				ctxInit: func(ctx *hinter.HintRunnerContext) {
+					err := ctx.ScopeManager.AssignVariable("keys", []fp.Element{})
+					if err != nil {
+						t.Fatal(err)
+					}
+				},
+				makeHinter: func(ctx *hintTestContext) hinter.Hinter {
+					return newSquashDictInnerNextKeyHint(ctx.operanders["next_key"])
+				},
+				errCheck: errorTextContains("no keys left but remaining_accesses > 0"),
+			},
+			{
+				operanders: []*hintOperander{
+					{Name: "next_key", Kind: uninitialized},
+				},
+				ctxInit: func(ctx *hinter.HintRunnerContext) {
+					err := ctx.ScopeManager.AssignVariable("keys", []fp.Element{*feltUint64(3), *feltUint64(2), *feltUint64(1)})
+					if err != nil {
+						t.Fatal(err)
+					}
+				},
+				makeHinter: func(ctx *hintTestContext) hinter.Hinter {
+					return newSquashDictInnerNextKeyHint(ctx.operanders["next_key"])
+				},
+				check: func(t *testing.T, ctx *hintTestContext) {
+					allVarValueInScopeEquals(map[string]any{"keys": []fp.Element{*feltUint64(3), *feltUint64(2)}, "key": *feltUint64((1))})(t, ctx)
+					varValueEquals("next_key", feltUint64(1))(t, ctx)
+				},
+			},
+			{
+				operanders: []*hintOperander{
+					{Name: "next_key", Kind: uninitialized},
+				},
+				ctxInit: func(ctx *hinter.HintRunnerContext) {
+					err := ctx.ScopeManager.AssignVariable("keys", []fp.Element{*feltUint64(15), *feltUint64(12), *feltUint64(9), *feltUint64(7), *feltUint64(6), *feltUint64(4)})
+					if err != nil {
+						t.Fatal(err)
+					}
+				},
+				makeHinter: func(ctx *hintTestContext) hinter.Hinter {
+					return newSquashDictInnerNextKeyHint(ctx.operanders["next_key"])
+				},
+				check: func(t *testing.T, ctx *hintTestContext) {
+					allVarValueInScopeEquals(map[string]any{"keys": []fp.Element{*feltUint64(15), *feltUint64(12), *feltUint64(9), *feltUint64(7), *feltUint64(6)}, "key": *feltUint64((4))})(t, ctx)
+					varValueEquals("next_key", feltUint64(4))(t, ctx)
+				},
+			},
+		},
 	})
 }
diff --git a/pkg/hintrunner/zero/zerohint_utils_test.go b/pkg/hintrunner/zero/zerohint_utils_test.go
@@ -13,6 +13,7 @@ import (
 	VM "github.com/NethermindEth/cairo-vm-go/pkg/vm"
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm/memory"
 	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
+	f "github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 	"github.com/stretchr/testify/require"
 )
 
@@ -172,10 +173,10 @@ func varValueInScopeEquals(varName string, expected any) func(t *testing.T, ctx
 					t.Fatalf("%s scope value mismatch:\nhave: %v\nwant: %v", varName, value, expected)
 				}
 			}
-		case []fp.Element:
+		case []f.Element:
 			{
-				valueArray := value.([]fp.Element)
-				expectedArray := expected.([]fp.Element)
+				valueArray := value.([]f.Element)
+				expectedArray := expected.([]f.Element)
 				if !reflect.DeepEqual(valueArray, expectedArray) {
 					t.Fatalf("%s scope value mismatch:\nhave: %v\nwant: %v", varName, value, expected)
 				}