Repository: NethermindEth/cairo-vm-go
Pull request #694 (2025-02-21 00:02:43)

Dict relocation


Patch:
diff --git a/cmd/cli/main.go b/cmd/cli/main.go
index 8f86f3f..4bc869a 100644
--- a/cmd/cli/main.go
+++ b/cmd/cli/main.go
@@ -287,12 +287,6 @@ func runVM(
 		if err := cairoRunner.FinalizeSegments(); err != nil {
 			return fmt.Errorf("cannot finalize segments: %w", err)
 		}
-	case runner.ExecutionModeCairo:
-		if airPublicInputLocation != "" {
-			if err := cairoRunner.FinalizeBuiltins(); err != nil {
-				return fmt.Errorf("cannot finalize builtins: %w", err)
-			}
-		}
 	case runner.ProofModeCairo:
 		if err := cairoRunner.EndRun(); err != nil {
 			return fmt.Errorf("cannot end run: %w", err)
@@ -301,9 +295,10 @@ func runVM(
 			if err := cairoRunner.FinalizeBuiltins(); err != nil {
 				return fmt.Errorf("cannot finalize builtins: %w", err)
 			}
-		}
-		if err := cairoRunner.FinalizeSegments(); err != nil {
-			return fmt.Errorf("cannot finalize segments: %w", err)
+
+			if err := cairoRunner.FinalizeSegments(); err != nil {
+				return fmt.Errorf("cannot finalize segments: %w", err)
+			}
 		}
 	}
 
@@ -319,6 +314,7 @@ func runVM(
 			}
 		}
 	}
+
 	var segmentsOffsets []uint64
 	var relocatedMemory []*fp.Element
 	if proofmode || buildMemory {
diff --git a/pkg/hintrunner/core/hint.go b/pkg/hintrunner/core/hint.go
index 584c8f5..5fe76d6 100644
--- a/pkg/hintrunner/core/hint.go
+++ b/pkg/hintrunner/core/hint.go
@@ -2021,9 +2021,6 @@ func (hint *RelocateAllDictionaries) String() string {
 }
 
 func (hint *RelocateAllDictionaries) Execute(vm *VM.VirtualMachine, ctx *hinter.HintRunnerContext) error {
-	err := ctx.DictionaryManager.RelocateAllDictionaries(vm)
-	if err != nil {
-		return fmt.Errorf("relocate all dictionaries: %w", err)
-	}
+	ctx.DictionaryManager.RelocateAllDictionaries(vm)
 	return nil
 }
diff --git a/pkg/hintrunner/hinter/dict.go b/pkg/hintrunner/hinter/dict.go
index ceb743d..894fb7c 100644
--- a/pkg/hintrunner/hinter/dict.go
+++ b/pkg/hintrunner/hinter/dict.go
@@ -14,6 +14,7 @@ type Dictionary struct {
 	data map[f.Element]*mem.MemoryValue
 	// Unique id assigned at the moment of creation
 	idx uint64
+	end mem.MemoryAddress
 }
 
 // Gets the memory value at certain key
@@ -29,6 +30,10 @@ func (d *Dictionary) Set(key *f.Element, value *mem.MemoryValue) {
 	d.data[*key] = value
 }
 
+func (d *Dictionary) SetEnd(end mem.MemoryAddress) {
+	d.end = end
+}
+
 // Returns the initialization number when the dictionary was created
 func (d *Dictionary) InitNumber() uint64 {
 	return d.idx
@@ -37,14 +42,14 @@ func (d *Dictionary) InitNumber() uint64 {
 // Used to manage dictionaries creation
 type DictionaryManager struct {
 	// a map that links a segment index to a dictionary
-	dictionaries map[uint64]Dictionary
+	dictionaries map[int]Dictionary
 	// useTemporarySegments is a flag that indicates if the dictionaries should be located in temporary segments, and later relocated to memory segments
 	useTemporarySegments bool
 }
 
 func InitializeDictionaryManager(ctx *HintRunnerContext, useTemporarySegments bool) {
 	if ctx.DictionaryManager.dictionaries == nil {
-		ctx.DictionaryManager.dictionaries = make(map[uint64]Dictionary)
+		ctx.DictionaryManager.dictionaries = make(map[int]Dictionary)
 	}
 	ctx.DictionaryManager.useTemporarySegments = useTemporarySegments
 }
@@ -89,6 +94,8 @@ func (dm *DictionaryManager) At(dictAddr *mem.MemoryAddress, key *f.Element) (*m
 func (dm *DictionaryManager) Set(dictAddr *mem.MemoryAddress, key *f.Element, value *mem.MemoryValue) error {
 	if dict, ok := dm.dictionaries[dictAddr.SegmentIndex]; ok {
 		dict.Set(key, value)
+		dict.SetEnd(*dictAddr)
+		dm.dictionaries[dictAddr.SegmentIndex] = dict
 		return nil
 	}
 	return fmt.Errorf("no dictionary at address %s", dictAddr)
@@ -96,31 +103,12 @@ func (dm *DictionaryManager) Set(dictAddr *mem.MemoryAddress, key *f.Element, va
 
 // Relocates all dictionaries into a single segment if proofmode is enabled
 // In LambdaClass VM there is add_relocation_rule() used, which is used only to relocate dictionaries / in specific hint. Thus we relocate dictionaries right away.
-func (dm *DictionaryManager) RelocateAllDictionaries(vm *VM.VirtualMachine) error {
+func (dm *DictionaryManager) RelocateAllDictionaries(vm *VM.VirtualMachine) {
 	segmentAddr := vm.Memory.AllocateEmptySegment()
 	for key, dict := range dm.dictionaries {
-		for _, value := range dict.data {
-			if value.IsAddress() {
-				addr, err := value.MemoryAddress()
-				if err != nil {
-					return err
-				}
-				element := key + addr.Offset
-				mv := mem.MemoryValueFromUint(element)
-				err = vm.Memory.WriteToAddress(&segmentAddr, &mv)
-				if err != nil {
-					return err
-				}
-			} else {
-				err := vm.Memory.WriteToAddress(&segmentAddr, value)
-				if err != nil {
-					return err
-				}
-			}
-			segmentAddr.Offset++
-		}
+		vm.Memory.AddRelocationRule(-key, segmentAddr)
+		segmentAddr.Offset += dict.end.Offset + 1
 	}
-	return nil
 }
 
 // Used to keep track of squashed dictionaries
diff --git a/pkg/hintrunner/hinter/zero_dict.go b/pkg/hintrunner/hinter/zero_dict.go
index d74cbb5..40d6935 100644
--- a/pkg/hintrunner/hinter/zero_dict.go
+++ b/pkg/hintrunner/hinter/zero_dict.go
@@ -46,12 +46,12 @@ func (d *ZeroDictionary) setFreeOffset(freeOffset uint64) {
 // Used to manage dictionaries creation
 type ZeroDictionaryManager struct {
 	// a map that links a segment index to a dictionary
-	Dictionaries map[uint64]ZeroDictionary
+	Dictionaries map[int]ZeroDictionary
 }
 
 func NewZeroDictionaryManager() ZeroDictionaryManager {
 	return ZeroDictionaryManager{
-		Dictionaries: make(map[uint64]ZeroDictionary),
+		Dictionaries: make(map[int]ZeroDictionary),
 	}
 }
 
diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go
index f684e34..67e2176 100644
--- a/pkg/runner/runner.go
+++ b/pkg/runner/runner.go
@@ -222,9 +222,6 @@ func (runner *Runner) initializeSegments() (*mem.Memory, error) {
 	}
 
 	memory.AllocateEmptySegment() // ExecutionSegment
-	if runner.isProofMode() {
-		memory.TemporarySegments = []*mem.Segment{}
-	}
 	return memory, nil
 }
 
@@ -444,6 +441,11 @@ func (runner *Runner) RunFor(steps uint64) error {
 // Since this vm always finishes the run of the program at the number of steps that is a power of two in the proof mode,
 // there is no need to run additional steps before the loop.
 func (runner *Runner) EndRun() error {
+	if runner.runnerMode == ProofModeCairo {
+		if err := runner.RelocateTemporarySegments(); err != nil {
+			return err
+		}
+	}
 	for runner.checkUsedCells() != nil {
 		pow2Steps := utils.NextPowerOfTwo(runner.vm.Step + 1)
 		if err := runner.RunFor(pow2Steps); err != nil {
@@ -604,6 +606,13 @@ func (runner *Runner) Output() []*fp.Element {
 	return output
 }
 
+func (runner *Runner) RelocateTemporarySegments() error {
+	if err := runner.vm.Memory.RelocateTemporarySegments(); err != nil {
+		return err
+	}
+	return nil
+}
+
 type InlineCasmContext struct {
 	instructions      []*fp.Element
 	currentCodeOffset int
@@ -661,15 +670,27 @@ func GetEntryCodeInstructions(function starknet.EntryPointByFunction, proofmode
 	}
 	hints := make(map[uint64][]hinter.Hinter)
 
+	paramsSize := 0
+	for _, param := range paramTypes {
+		paramsSize += param.Size
+	}
+
+	// The hint can be executed before the first instruction, because the AP correction was calculated based on the input arguments.
+	if paramsSize > 0 {
+		hints[uint64(0)] = append(hints[uint64(0)], []hinter.Hinter{
+			&core.ExternalWriteArgsToMemory{},
+		}...)
+	}
+
 	if gotSegmentArena {
-		hints[uint64(ctx.currentCodeOffset)] = []hinter.Hinter{
+		hints[uint64(ctx.currentCodeOffset)] = append(hints[uint64(ctx.currentCodeOffset)], []hinter.Hinter{
 			&core.AllocSegment{
 				Dst: hinter.ApCellRef(0),
 			},
 			&core.AllocSegment{
 				Dst: hinter.ApCellRef(1),
 			},
-		}
+		}...)
 		ctx.AddInlineCASM(
 			"[ap+2] = 0, ap++;",
 		)
@@ -685,10 +706,6 @@ func GetEntryCodeInstructions(function starknet.EntryPointByFunction, proofmode
 		apOffset += 3
 	}
 
-	paramsSize := 0
-	for _, param := range paramTypes {
-		paramsSize += param.Size
-	}
 	apOffset += paramsSize
 	gotGasBuiltin := false
 
@@ -719,13 +736,6 @@ func GetEntryCodeInstructions(function starknet.EntryPointByFunction, proofmode
 		)
 	}
 
-	// The hint can be executed before the first instruction, because the AP correction was calculated based on the input arguments.
-	if paramsSize > 0 {
-		hints[uint64(0)] = append(hints[uint64(0)], []hinter.Hinter{
-			&core.ExternalWriteArgsToMemory{},
-		}...)
-	}
-
 	codeOffsetBeforeCallRel := uint64(codeOffset) - uint64(ctx.currentCodeOffset)
 	ctx.AddInlineCASM("call rel 0;")
 	callRelArgLocation := len(ctx.instructions) - 1
diff --git a/pkg/vm/memory/memory.go b/pkg/vm/memory/memory.go
index b654790..79aec43 100644
--- a/pkg/vm/memory/memory.go
+++ b/pkg/vm/memory/memory.go
@@ -215,6 +215,7 @@ type Memory struct {
 	Segments []*Segment
 	// TemporarySegments is a map of temporary segments, key is the segment index, value is the segment
 	TemporarySegments []*Segment
+	relocationRules   map[int]MemoryAddress
 }
 
 // todo(rodro): can the amount of segments be known before hand?
@@ -222,6 +223,9 @@ func InitializeEmptyMemory() *Memory {
 	return &Memory{
 		// capacity 4 should be enough for the minimum amount of segments
 		Segments: make([]*Segment, 0, 4),
+		// allocate 1 empty temporary segment, so that the real first segment index is 1, indexed by -1 in temporary address
+		TemporarySegments: []*Segment{EmptySegment()},
+		relocationRules:   make(map[int]MemoryAddress),
 	}
 }
 
@@ -237,7 +241,7 @@ func (memory *Memory) AllocateSegment(data []*f.Element) (MemoryAddress, error)
 	}
 	memory.Segments = append(memory.Segments, newSegment)
 	return MemoryAddress{
-		SegmentIndex: uint64(len(memory.Segments) - 1),
+		SegmentIndex: len(memory.Segments) - 1,
 		Offset:       0,
 	}, nil
 }
@@ -246,7 +250,7 @@ func (memory *Memory) AllocateSegment(data []*f.Element) (MemoryAddress, error)
 func (memory *Memory) AllocateEmptySegment() MemoryAddress {
 	memory.Segments = append(memory.Segments, EmptySegment())
 	return MemoryAddress{
-		SegmentIndex: uint64(len(memory.Segments) - 1),
+		SegmentIndex: len(memory.Segments) - 1,
 		Offset:       0,
 	}
 }
@@ -255,7 +259,7 @@ func (memory *Memory) AllocateEmptySegment() MemoryAddress {
 func (memory *Memory) AllocateEmptyTemporarySegment() MemoryAddress {
 	memory.TemporarySegments = append(memory.TemporarySegments, EmptySegment())
 	return MemoryAddress{
-		SegmentIndex: uint64(len(memory.TemporarySegments)),
+		SegmentIndex: -(len(memory.TemporarySegments) - 1),
 		Offset:       0,
 	}
 }
@@ -265,21 +269,32 @@ func (memory *Memory) AllocateBuiltinSegment(builtinRunner BuiltinRunner) Memory
 	builtinSegment := EmptySegment().WithBuiltinRunner(builtinRunner)
 	memory.Segments = append(memory.Segments, builtinSegment)
 	return MemoryAddress{
-		SegmentIndex: uint64(len(memory.Segments) - 1),
+		SegmentIndex: len(memory.Segments) - 1,
 		Offset:       0,
 	}
 }
 
 // Writes to a given segment index and offset a new memory value. Errors if writing
 // to an unallocated segment or if overwriting a different memory value
-func (memory *Memory) Write(segmentIndex uint64, offset uint64, value *MemoryValue) error {
-	if segmentIndex >= uint64(len(memory.Segments)) {
-		return fmt.Errorf("segment %d: unallocated", segmentIndex)
-	}
-	if err := memory.Segments[segmentIndex].Write(offset, value); err != nil {
-		return fmt.Errorf("segment %d, offset %d: %w", segmentIndex, offset, err)
+func (memory *Memory) Write(segmentIndex int, offset uint64, value *MemoryValue) error {
+	if segmentIndex >= 0 {
+		if segmentIndex >= len(memory.Segments) {
+			return fmt.Errorf("segment %d: unallocated", segmentIndex)
+		}
+		if err := memory.Segments[segmentIndex].Write(offset, value); err != nil {
+			return fmt.Errorf("segment %d, offset %d: %w", segmentIndex, offset, err)
+		}
+		return nil
+	} else {
+		segmentIndex = -segmentIndex
+		if segmentIndex >= len(memory.TemporarySegments) {
+			return fmt.Errorf("temporary segment %d: unallocated", segmentIndex)
+		}
+		if err := memory.TemporarySegments[segmentIndex].Write(offset, value); err != nil {
+			return fmt.Errorf("temporary segment %d, offset %d: %w", segmentIndex, offset, err)
+		}
+		return nil
 	}
-	return nil
 }
 
 // Writes to a memory address a new memory value. Errors if writing to an unallocated
@@ -290,15 +305,27 @@ func (memory *Memory) WriteToAddress(address *MemoryAddress, value *MemoryValue)
 
 // Reads a memory value given the segment index and offset. Errors if reading from
 // an unallocated segment or if reading an unknown memory value
-func (memory *Memory) Read(segmentIndex uint64, offset uint64) (MemoryValue, error) {
-	if segmentIndex >= uint64(len(memory.Segments)) {
-		return MemoryValue{}, fmt.Errorf("segment %d: unallocated", segmentIndex)
-	}
-	mv, err := memory.Segments[segmentIndex].Read(offset)
-	if err != nil {
-		return MemoryValue{}, fmt.Errorf("segment %d, offset %d: %w", segmentIndex, offset, err)
+func (memory *Memory) Read(segmentIndex int, offset uint64) (MemoryValue, error) {
+	if segmentIndex >= 0 {
+		if segmentIndex >= len(memory.Segments) {
+			return MemoryValue{}, fmt.Errorf("segment %d: unallocated", segmentIndex)
+		}
+		mv, err := memory.Segments[segmentIndex].Read(offset)
+		if err != nil {
+			return MemoryValue{}, fmt.Errorf("segment %d, offset %d: %w", segmentIndex, offset, err)
+		}
+		return mv, nil
+	} else {
+		segmentIndex = -segmentIndex
+		if segmentIndex >= len(memory.TemporarySegments) {
+			return MemoryValue{}, fmt.Errorf("temporary segment %d: unallocated", segmentIndex)
+		}
+		mv, err := memory.TemporarySegments[segmentIndex].Read(offset)
+		if err != nil {
+			return MemoryValue{}, fmt.Errorf("temporary segment %d, offset %d: %w", segmentIndex, offset, err)
+		}
+		return mv, nil
 	}
-	return mv, nil
 }
 
 // Reads a memory value given an address. Errors if reading from
@@ -308,7 +335,7 @@ func (memory *Memory) ReadFromAddress(address *MemoryAddress) (MemoryValue, erro
 }
 
 // Works the same as `Read` but `MemoryValue` is converted to `Element` first
-func (memory *Memory) ReadAsElement(segmentIndex uint64, offset uint64) (f.Element, error) {
+func (memory *Memory) ReadAsElement(segmentIndex int, offset uint64) (f.Element, error) {
 	mv, err := memory.Read(segmentIndex, offset)
 	if err != nil {
 		return f.Element{}, err
@@ -345,11 +372,19 @@ func (memory *Memory) ReadFromAddressAsAddress(address *MemoryAddress) (MemoryAd
 
 // Given a segment index and offset, returns the memory value at that position, without
 // modifying it in any way. Errors if peeking from an unallocated segment
-func (memory *Memory) Peek(segmentIndex uint64, offset uint64) (MemoryValue, error) {
-	if segmentIndex >= uint64(len(memory.Segments)) {
-		return MemoryValue{}, fmt.Errorf("segment %d: unallocated", segmentIndex)
+func (memory *Memory) Peek(segmentIndex int, offset uint64) (MemoryValue, error) {
+	if segmentIndex >= 0 {
+		if segmentIndex >= len(memory.Segments) {
+			return MemoryValue{}, fmt.Errorf("segment %d: unallocated", segmentIndex)
+		}
+		return memory.Segments[segmentIndex].Peek(offset), nil
+	} else {
+		segmentIndex = -segmentIndex
+		if segmentIndex >= len(memory.TemporarySegments) {
+			return MemoryValue{}, fmt.Errorf("temporary segment %d: unallocated", segmentIndex)
+		}
+		return memory.TemporarySegments[segmentIndex].Peek(offset), nil
 	}
-	return memory.Segments[segmentIndex].Peek(offset), nil
 }
 
 // Given an address returns the memory value at that position, without
@@ -360,12 +395,21 @@ func (memory *Memory) PeekFromAddress(address *MemoryAddress) (MemoryValue, erro
 
 // Given a segment index and offset returns true if the value at that address
 // is known
-func (memory *Memory) KnownValue(segment uint64, offset uint64) bool {
-	if segment >= uint64(len(memory.Segments)) ||
-		offset >= uint64(len(memory.Segments[segment].Data)) {
-		return false
+func (memory *Memory) KnownValue(segment int, offset uint64) bool {
+	if segment >= 0 {
+		if segment >= len(memory.Segments) ||
+			offset >= uint64(len(memory.Segments[segment].Data)) {
+			return false
+		}
+		return memory.Segments[segment].Data[offset].Known()
+	} else {
+		segment = -segment
+		if segment >= len(memory.TemporarySegments) ||
+			offset >= uint64(len(memory.TemporarySegments[segment].Data)) {
+			return false
+		}
+		return memory.TemporarySegments[segment].Data[offset].Known()
 	}
-	return memory.Segments[segment].Data[offset].Known()
 }
 
 // Given an address returns true if it contains a known value
@@ -417,3 +461,51 @@ func (memory *Memory) WriteToNthStructField(addr MemoryAddress, value MemoryValu
 	}
 	return memory.WriteToAddress(&nAddr, &value)
 }
+
+func (memory *Memory) AddRelocationRule(segmentIndex int, addr MemoryAddress) {
+	memory.relocationRules[segmentIndex] = addr
+}
+
+func (memory *Memory) RelocateTemporarySegments() error {
+	// We check if the length of the temporary segments is 1 because the first temporary is added during initialization
+	// for proper indexing, and is always empty
+	if len(memory.relocationRules) == 0 || len(memory.TemporarySegments)-1 == 0 {
+		return nil
+	}
+	for i, segment := range memory.Segments {
+		for j := uint64(0); j < segment.RealLen(); j++ {
+			if !segment.Data[j].Known() {
+				continue
+			}
+
+			if segment.Data[j].IsAddress() {
+				addr, _ := segment.Data[j].MemoryAddress()
+				if addr.SegmentIndex < 0 {
+					if rule, ok := memory.relocationRules[-addr.SegmentIndex]; ok {
+						newAddr := MemoryAddress{SegmentIndex: rule.SegmentIndex, Offset: rule.Offset + addr.Offset}
+						memory.Segments[i].Data[j] = MemoryValueFromMemoryAddress(&newAddr)
+					}
+				}
+			}
+		}
+	}
+
+	for index := 0; index < len(memory.TemporarySegments); index++ {
+		baseAddr, ok := memory.relocationRules[index]
+		if !ok {
+			continue
+		}
+
+		dataSegment := memory.TemporarySegments[index]
+
+		for _, cell := range dataSegment.Data {
+			if cell.Known() {
+				if err := memory.Write(baseAddr.SegmentIndex, baseAddr.Offset, &cell); err != nil {
+					return err
+				}
+				baseAddr.Offset++
+			}
+		}
+	}
+	return nil
+}
diff --git a/pkg/vm/memory/memory_value.go b/pkg/vm/memory/memory_value.go
index ebde1d8..d8964ba 100644
--- a/pkg/vm/memory/memory_value.go
+++ b/pkg/vm/memory/memory_value.go
@@ -14,7 +14,7 @@ import (
 // during execution, addresses has two locators: the segment they belong to and the location
 // inside that segment
 type MemoryAddress struct {
-	SegmentIndex uint64
+	SegmentIndex int
 	Offset       uint64
 }
 
@@ -162,7 +162,7 @@ func MemoryValueFromUint[T constraints.Unsigned](v T) MemoryValue {
 func MemoryValueFromSegmentAndOffset[T constraints.Integer](segmentIndex, offset T) MemoryValue {
 	return MemoryValueFromMemoryAddress(
 		&MemoryAddress{
-			SegmentIndex: uint64(segmentIndex),
+			SegmentIndex: int(segmentIndex),
 			Offset:       uint64(offset),
 		},
 	)
diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 4c79428..69e91d9 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -92,8 +92,11 @@ type VirtualMachine struct {
 	RcLimitsMax uint16
 }
 
-func (vm *VirtualMachine) PrintMemory() {
+func (vm *VirtualMachine) PrintMemory(skipBytecode bool) {
 	for i := range vm.Memory.Segments {
+		if skipBytecode && i == ProgramSegment {
+			continue
+		}
 		for j, cell := range vm.Memory.Segments[i].Data {
 			if !cell.Known() {
 				continue
@@ -412,7 +415,7 @@ func (vm *VirtualMachine) opcodeAssertions(
 
 		apMv := mem.MemoryValueFromSegmentAndOffset(
 			vm.Context.Pc.SegmentIndex,
-			vm.Context.Pc.Offset+uint64(instruction.Size()),
+			int(vm.Context.Pc.Offset+uint64(instruction.Size())),
 		)
 		// Write in [ap + 1] the next instruction to execute
 		if err := vm.Memory.WriteToAddress(op0Addr, &apMv); err != nil {
@@ -701,7 +704,7 @@ func (vm *VirtualMachine) BuiltinsFinalStackFromStackPointerDict(builtinNameToSt
 		if err != nil {
 			return err
 		}
-		if stop_pointer.SegmentIndex != uint64(segmentIndex) {
+		if stop_pointer.SegmentIndex != segmentIndex {
 			return fmt.Errorf("stop pointer segment index mismatch")
 		}
 		stopPointerOffset := stop_pointer.Offset


Test patch:
diff --git a/pkg/hintrunner/utils/testutils.go b/pkg/hintrunner/utils/testutils.go
index 9cf8bf3..6994905 100644
--- a/pkg/hintrunner/utils/testutils.go
+++ b/pkg/hintrunner/utils/testutils.go
@@ -5,14 +5,14 @@ import (
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm/memory"
 )
 
-func WriteTo(vm *VM.VirtualMachine, segment uint64, offset uint64, val memory.MemoryValue) {
+func WriteTo(vm *VM.VirtualMachine, segment int, offset uint64, val memory.MemoryValue) {
 	err := vm.Memory.Write(segment, offset, &val)
 	if err != nil {
 		panic(err)
 	}
 }
 
-func ReadFrom(vm *VM.VirtualMachine, segment uint64, offset uint64) memory.MemoryValue {
+func ReadFrom(vm *VM.VirtualMachine, segment int, offset uint64) memory.MemoryValue {
 	val, err := vm.Memory.Read(segment, offset)
 	if err != nil {
 		panic(err)
diff --git a/pkg/hintrunner/zero/zerohint_utils_test.go b/pkg/hintrunner/zero/zerohint_utils_test.go
index d817f7a..f02f7d5 100644
--- a/pkg/hintrunner/zero/zerohint_utils_test.go
+++ b/pkg/hintrunner/zero/zerohint_utils_test.go
@@ -23,7 +23,7 @@ func addr(offset uint64) *memory.MemoryAddress {
 	}
 }
 
-func addrWithSegment(segment, offset uint64) *memory.MemoryAddress {
+func addrWithSegment(segment int, offset uint64) *memory.MemoryAddress {
 	return &memory.MemoryAddress{
 		SegmentIndex: segment,
 		Offset:       offset,
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index 43bfe73..eaaa0fc 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -1366,13 +1366,13 @@ func newElementPtr(val uint64) *f.Element {
 }
 
 type memoryWrite struct {
-	SegmentIndex uint64
+	SegmentIndex int
 	Offset       uint64
 	Value        any
 }
 
 func updateMemoryWithValues(memory *mem.Memory, valuesToWrite []memoryWrite) {
-	var max_segment uint64 = 0
+	var max_segment int = 0
 	for _, toWrite := range valuesToWrite {
 		// wrap any inside a memory value
 		val, err := mem.MemoryValueFromAny(toWrite.Value)