Implement DictSquashUpdatePtr

pkg/hintrunner/hinter/zero_dict.go

@@ -10,7 +10,7 @@ import (
 
 // Used to keep track of all dictionaries data
 type ZeroDictionary struct {
-	// The Data contained on a dictionary
+	// The Data contained in a dictionary
 	Data map[f.Element]mem.MemoryValue
 	// Default value for key not present in the dictionary
 	DefaultValue mem.MemoryValue
@@ -19,7 +19,7 @@ type ZeroDictionary struct {
 }
 
 // Gets the memory value at certain key
-func (d *ZeroDictionary) At(key f.Element) (mem.MemoryValue, error) {
+func (d *ZeroDictionary) at(key f.Element) (mem.MemoryValue, error) {
 	if value, ok := d.Data[key]; ok {
 		return value, nil
 	}
@@ -30,15 +30,20 @@ func (d *ZeroDictionary) At(key f.Element) (mem.MemoryValue, error) {
 }
 
 // Given a key and a value, it sets the value at the given key
-func (d *ZeroDictionary) Set(key f.Element, value mem.MemoryValue) {
+func (d *ZeroDictionary) set(key f.Element, value mem.MemoryValue) {
 	d.Data[key] = value
 }
 
 // Given a incrementBy value, it increments the freeOffset field of dictionary by it
-func (d *ZeroDictionary) IncrementFreeOffset(freeOffset uint64) {
+func (d *ZeroDictionary) incrementFreeOffset(freeOffset uint64) {
 	*d.FreeOffset += freeOffset
 }
 
+// Given a freeOffset value, it sets the freeOffset field of dictionary to it
+func (d *ZeroDictionary) setFreeOffset(freeOffset uint64) {
+	*d.FreeOffset = freeOffset
+}
+
 // Used to manage dictionaries creation
 type ZeroDictionaryManager struct {
 	// a map that links a segment index to a dictionary
@@ -94,24 +99,33 @@ func (dm *ZeroDictionaryManager) GetDictionary(dictAddr mem.MemoryAddress) (Zero
 // to locate the correct dictionary and the key to index on it
 func (dm *ZeroDictionaryManager) At(dictAddr mem.MemoryAddress, key f.Element) (mem.MemoryValue, error) {
 	if dict, ok := dm.dictionaries[dictAddr.SegmentIndex]; ok {
-		return dict.At(key)
+		return dict.at(key)
 	}
 	return mem.UnknownValue, fmt.Errorf("no dictionary at address: %s", dictAddr)
 }
 
 // Given a memory address,a key and a value it stores the value at the correct position.
 func (dm *ZeroDictionaryManager) Set(dictAddr mem.MemoryAddress, key f.Element, value mem.MemoryValue) error {
 	if dict, ok := dm.dictionaries[dictAddr.SegmentIndex]; ok {
-		dict.Set(key, value)
+		dict.set(key, value)
 		return nil
 	}
 	return fmt.Errorf("no dictionary at address: %s", dictAddr)
 }
 
 // Given a memory address and a incrementBy, it increments the freeOffset field of dictionary by it.
-func (dm *ZeroDictionaryManager) IncrementFreeOffset(dictAddr mem.MemoryAddress, freeOffset uint64) error {
+func (dm *ZeroDictionaryManager) IncrementFreeOffset(dictAddr mem.MemoryAddress, incrementBy uint64) error {
+	if dict, ok := dm.dictionaries[dictAddr.SegmentIndex]; ok {
+		dict.incrementFreeOffset(incrementBy)
+		return nil
+	}
+	return fmt.Errorf("no dictionary at address: %s", dictAddr)
+}
+
+// Given a memory address and a freeOffset, it sets the freeOffset field of dictionary to it.
+func (dm *ZeroDictionaryManager) SetFreeOffset(dictAddr mem.MemoryAddress, freeOffset uint64) error {
 	if dict, ok := dm.dictionaries[dictAddr.SegmentIndex]; ok {
-		dict.IncrementFreeOffset(freeOffset)
+		dict.setFreeOffset(freeOffset)
 		return nil
 	}
 	return fmt.Errorf("no dictionary at address: %s", dictAddr)

pkg/hintrunner/zero/hintcode.go

@@ -139,6 +139,7 @@ const (
 	squashDictInnerLenAssert          string = "assert len(current_access_indices) == 0"
 	squashDictInnerNextKey            string = "assert len(keys) > 0, 'No keys left but remaining_accesses > 0.'\nids.next_key = key = keys.pop()"
 	squashDictInnerUsedAccessesAssert string = "assert ids.n_used_accesses == len(access_indices[key])"
+	dictSquashUpdatePtrCode           string = "# Update the DictTracker's current_ptr to point to the end of the squashed dict.\n__dict_manager.get_tracker(ids.squashed_dict_start).current_ptr = \\\n    ids.squashed_dict_end.address_"
 
 	// ------ Other hints related code ------
 	allocSegmentCode          string = "memory[ap] = segments.add()"

pkg/hintrunner/zero/zerohint.go

@@ -201,6 +201,8 @@ func GetHintFromCode(program *zero.ZeroProgram, rawHint zero.Hint, hintPC uint64
 		return createSquashDictInnerNextKeyHinter(resolver)
 	case squashDictInnerUsedAccessesAssert:
 		return createSquashDictInnerUsedAccessesAssertHinter(resolver)
+	case dictSquashUpdatePtrCode:
+		return createDictSquashUpdatePtrHinter(resolver)
 	// Other hints
 	case allocSegmentCode:
 		return createAllocSegmentHinter()

pkg/hintrunner/zero/zerohint_dictionaries.go

@@ -816,3 +816,47 @@ func createSquashDictInnerUsedAccessesAssertHinter(resolver hintReferenceResolve
 
 	return newSquashDictInnerUsedAccessesAssertHint(nUsedAccesses), nil
 }
+
+// DictSquashUpdatePtr updates the DictTracker's current_ptr to point to the end of the squashed dict
+//
+// `newDictSquashUpdatePtrHint` takes two operanders as arguments
+//   - `squashed_dict_start` pointer to the dictionary whose current_ptr should be updated
+//   - `squashed_dict_end` new current_ptr of the dictionary
+func newDictSquashUpdatePtrHint(squashedDictStart, squashedDictEnd hinter.ResOperander) hinter.Hinter {
+	return &GenericZeroHinter{
+		Name: "DictSquashUpdatePtr",
+		Op: func(vm *VM.VirtualMachine, ctx *hinter.HintRunnerContext) error {
+			//> __dict_manager.get_tracker(ids.squashed_dict_start).current_ptr = ids.squashed_dict_end.address_
+
+			squashedDictStart, err := hinter.ResolveAsAddress(vm, squashedDictStart)
+			if err != nil {
+				return err
+			}
+			squashedDictEnd, err := hinter.ResolveAsAddress(vm, squashedDictEnd)
+			if err != nil {
+				return err
+			}
+
+			dictionaryManager, ok := ctx.ScopeManager.GetZeroDictionaryManager()
+			if !ok {
+				return fmt.Errorf("__dict_manager not in scope")
+			}
+
+			// TODO: figure out if its ever possible for squashedDictEnd segment to be different from the dictionary segment
+			return dictionaryManager.SetFreeOffset(*squashedDictStart, squashedDictEnd.Offset)
+		},
+	}
+}
+
+func createDictSquashUpdatePtrHinter(resolver hintReferenceResolver) (hinter.Hinter, error) {
+	squashedDictStart, err := resolver.GetResOperander("squashed_dict_start")
+	if err != nil {
+		return nil, err
+	}
+	squashedDictEnd, err := resolver.GetResOperander("squashed_dict_end")
+	if err != nil {
+		return nil, err
+	}
+
+	return newDictSquashUpdatePtrHint(squashedDictStart, squashedDictEnd), nil
+}

pkg/hintrunner/zero/zerohint_dictionaries_test.go

@@ -49,13 +49,9 @@ func TestZeroHintDictionaries(t *testing.T) {
 						t.Fatalf("__dict_manager missing")
 					}
 					dictionaryManager := dictionaryManagerValue.(hinter.ZeroDictionaryManager)
-					dictionary, err := dictionaryManager.GetDictionary(dictAddr)
-					if err != nil {
-						t.Fatalf("error fetching dictionary from address at ap")
-					}
 
 					for _, key := range []fp.Element{*feltUint64(10), *feltUint64(20), *feltUint64(30)} {
-						value, err := dictionary.At(key)
+						value, err := dictionaryManager.At(dictAddr, key)
 						if err != nil {
 							t.Fatalf("error fetching value for key: %v", key)
 						}
@@ -577,5 +573,34 @@ func TestZeroHintDictionaries(t *testing.T) {
 				errCheck: errorTextContains("assertion ids.n_used_accesses == len(access_indices[key]) failed"),
 			},
 		},
+		"DictSquashUpdatePtr": {
+			{
+				operanders: []*hintOperander{
+					{Name: "squashed_dict_start", Kind: apRelative, Value: addrWithSegment(2, 0)},
+					{Name: "squashed_dict_end", Kind: apRelative, Value: addrWithSegment(2, 8)},
+				},
+				makeHinter: func(ctx *hintTestContext) hinter.Hinter {
+					dictionaryManager := hinter.NewZeroDictionaryManager()
+					err := ctx.runnerContext.ScopeManager.AssignVariable("__dict_manager", dictionaryManager)
+					if err != nil {
+						t.Fatal(err)
+					}
+					defaultValueMv := memory.MemoryValueFromInt(12345)
+					dictionaryManager.NewDefaultDictionary(ctx.vm, defaultValueMv)
+
+					return newDictSquashUpdatePtrHint(
+						ctx.operanders["squashed_dict_start"],
+						ctx.operanders["squashed_dict_end"],
+					)
+				},
+				check: func(t *testing.T, ctx *hintTestContext) {
+					dictPtr := addrWithSegment(2, 0)
+					expectedData := map[fp.Element]memory.MemoryValue{}
+					expectedDefaultValue := memory.MemoryValueFromInt(12345)
+					expectedFreeOffset := uint64(8)
+					zeroDictInScopeEquals(*dictPtr, expectedData, expectedDefaultValue, expectedFreeOffset)(t, ctx)
+				},
+			},
+		},
 	})
 }