Memory relocation

pkg/vm/memory/memory.go

@@ -15,6 +15,10 @@ func NewMemory() *Memory {
 	return &Memory{data, 0}
 }
 
+func (m *Memory) NumSegments() uint {
+	return m.num_segments
+}
+
 // Inserts a value in some memory address, given by a Relocatable value.
 func (m *Memory) Insert(addr Relocatable, val *MaybeRelocatable) error {
 	// FIXME: There should be a special handling if the key

pkg/vm/memory/relocatable.go

@@ -1,5 +1,10 @@
 package memory
 
+import (
+	"errors"
+	"fmt"
+)
+
 // Relocatable in the Cairo VM represents an address
 // in some memory segment. When the VM finishes running,
 // these values are replaced by real memory addresses,
@@ -15,6 +20,10 @@ func NewRelocatable(segment_idx int, offset uint) Relocatable {
 	return Relocatable{segment_idx, offset}
 }
 
+func (r *Relocatable) RelocateAddress(relocationTable *[]uint) uint {
+	return (*relocationTable)[r.SegmentIndex] + r.Offset
+}
+
 // Int in the Cairo VM represents a value in memory that
 // is not an address.
 type Int struct {
@@ -52,3 +61,21 @@ func (m *MaybeRelocatable) GetRelocatable() (Relocatable, bool) {
 	rel, is_type := m.inner.(Relocatable)
 	return rel, is_type
 }
+
+// Turns a MaybeRelocatable into a Felt252 value.
+// If the inner value is an Int, it will extract the Felt252 value from it.
+// If the inner value is a Relocatable, it will relocate it according to the relocation_table
+// TODO: Return value should be of type (felt, error)
+func (m *MaybeRelocatable) RelocateValue(relocationTable *[]uint) (uint, error) {
+	inner_int, ok := m.GetInt()
+	if ok {
+		return inner_int.felt, nil
+	}
+
+	inner_relocatable, ok := m.GetRelocatable()
+	if ok {
+		return inner_relocatable.RelocateAddress(relocationTable), nil
+	}
+
+	return 0, errors.New(fmt.Sprintf("Unexpected type %T", m.inner))
+}

pkg/vm/memory/segments.go

@@ -4,7 +4,7 @@ package memory
 // Also holds metadata useful for the relocation process of
 // the memory at the end of the VM run.
 type MemorySegmentManager struct {
-	segmentSizes map[uint]uint
+	SegmentSizes map[uint]uint
 	Memory       Memory
 }
 
@@ -20,6 +20,61 @@ func (m *MemorySegmentManager) AddSegment() Relocatable {
 	return ptr
 }
 
+// Calculates the size of each memory segment.
+func (m *MemorySegmentManager) ComputeEffectiveSizes() map[uint]uint {
+	if len(m.SegmentSizes) == 0 {
+
+		for ptr := range m.Memory.data {
+			segmentIndex := uint(ptr.SegmentIndex)
+			segmentMaxSize := m.SegmentSizes[segmentIndex]
+			segmentSize := ptr.Offset + 1
+			if segmentSize > segmentMaxSize {
+				m.SegmentSizes[segmentIndex] = segmentSize
+			}
+		}
+	}
+
+	return m.SegmentSizes
+}
+
+func (m *MemorySegmentManager) RelocateSegments() ([]uint, bool) {
+	if m.SegmentSizes == nil {
+		return nil, false
+	}
+
+	first_addr := uint(1)
+	relocation_table := []uint{first_addr}
+
+	for i := uint(0); i < m.Memory.NumSegments(); i++ {
+		new_addr := relocation_table[i] + m.SegmentSizes[i]
+		relocation_table = append(relocation_table, new_addr)
+	}
+	relocation_table = relocation_table[:len(relocation_table)-1]
+
+	return relocation_table, true
+}
+
+func (s *MemorySegmentManager) RelocateMemory(relocationTable *[]uint) (map[uint]uint, error) {
+	relocatedMemory := make(map[uint]uint, 0)
+
+	for i := uint(0); i < s.Memory.NumSegments(); i++ {
+		for j := uint(0); j < s.SegmentSizes[i]; j++ {
+			ptr := NewRelocatable(int(i), j)
+			cell, err := s.Memory.Get(ptr)
+			if err == nil {
+				relocatedAddr := ptr.RelocateAddress(relocationTable)
+				value, err := cell.RelocateValue(relocationTable)
+				if err != nil {
+					return nil, err
+				}
+				relocatedMemory[relocatedAddr] = value
+			}
+		}
+	}
+
+	return relocatedMemory, nil
+}
+
 // Writes data into the memory from address ptr and returns the first address after the data.
 // If any insertion fails, returns (0,0) and the memory insertion error
 func (m *MemorySegmentManager) LoadData(ptr Relocatable, data *[]MaybeRelocatable) (Relocatable, error) {

pkg/vm/memory/segments_test.go

@@ -0,0 +1,216 @@
+package memory_test
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/lambdaclass/cairo-vm.go/pkg/vm"
+	"github.com/lambdaclass/cairo-vm.go/pkg/vm/memory"
+)
+
+func TestComputeEffectiveSizeOneSegment(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 0), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 1), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 2), memory.NewMaybeRelocatableInt(1))
+
+	segments.ComputeEffectiveSizes()
+
+	expectedSizes := map[uint]uint{0: 3}
+	if !reflect.DeepEqual(expectedSizes, segments.SegmentSizes) {
+		t.Errorf("Segment sizes are not the same")
+	}
+}
+
+func TestComputeEffectiveSizeOneSegmentWithOneGap(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 6), memory.NewMaybeRelocatableInt(1))
+
+	segments.ComputeEffectiveSizes()
+
+	expectedSizes := map[uint]uint{0: 7}
+	if !reflect.DeepEqual(expectedSizes, segments.SegmentSizes) {
+		t.Errorf("Segment sizes are not the same")
+	}
+}
+
+func TestComputeEffectiveSizeOneSegmentWithMultipleGaps(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 3), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 4), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 7), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 9), memory.NewMaybeRelocatableInt(1))
+
+	segments.ComputeEffectiveSizes()
+
+	expectedSizes := map[uint]uint{0: 10}
+	if !reflect.DeepEqual(expectedSizes, segments.SegmentSizes) {
+		t.Errorf("Segment sizes are not the same")
+	}
+}
+
+func TestComputeEffectiveSizeThreeSegments(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 0), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 1), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 2), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(1, 0), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(1, 1), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(1, 2), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 0), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 1), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 2), memory.NewMaybeRelocatableInt(1))
+
+	segments.ComputeEffectiveSizes()
+
+	expectedSizes := map[uint]uint{0: 3, 1: 3, 2: 3}
+	if !reflect.DeepEqual(expectedSizes, segments.SegmentSizes) {
+		t.Errorf("Segment sizes are not the same")
+	}
+}
+
+func TestComputeEffectiveSizeThreeSegmentsWithGaps(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 2), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 5), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 7), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(1, 1), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 2), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 4), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 7), memory.NewMaybeRelocatableInt(1))
+
+	segments.ComputeEffectiveSizes()
+
+	expectedSizes := map[uint]uint{0: 8, 1: 2, 2: 8}
+	if !reflect.DeepEqual(expectedSizes, segments.SegmentSizes) {
+		t.Errorf("Segment sizes are not the same")
+	}
+}
+
+func TestGetSegmentUsedSizeAfterComputingUsed(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.Memory.Insert(memory.NewRelocatable(0, 2), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 5), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(0, 7), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(1, 1), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 2), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 4), memory.NewMaybeRelocatableInt(1))
+	segments.Memory.Insert(memory.NewRelocatable(2, 7), memory.NewMaybeRelocatableInt(1))
+
+	segments.ComputeEffectiveSizes()
+
+	segmentSize, ok := segments.SegmentSizes[2]
+	expectedSize := 8
+	if !ok || segmentSize != uint(expectedSize) {
+		t.Errorf("Segment size should be %d but it's %d", expectedSize, segmentSize)
+	}
+}
+
+func TestGetSegmentUsedSizeBeforeComputingUsed(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+
+	_, ok := segments.SegmentSizes[2]
+	if ok {
+		t.Errorf("Expected no segment sizes loaded")
+	}
+}
+
+func TestRelocateOneSegment(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.SegmentSizes = map[uint]uint{0: 3}
+	relocationTable, ok := segments.RelocateSegments()
+
+	if !ok {
+		t.Errorf("Memory segment manager doesn't have segment sizes initialized")
+	}
+
+	expectedTable := []uint{1}
+	if !reflect.DeepEqual(expectedTable, relocationTable) {
+		t.Errorf("Relocation tables are not the same")
+	}
+}
+
+func TestRelocateFiveSegments(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.SegmentSizes = map[uint]uint{0: 3, 1: 3, 2: 56, 3: 78, 4: 8}
+	relocationTable, ok := segments.RelocateSegments()
+
+	if !ok {
+		t.Errorf("Memory segment manager doesn't have segment sizes initialized")
+	}
+
+	expectedTable := []uint{1, 4, 7, 63, 141}
+	if !reflect.DeepEqual(expectedTable, relocationTable) {
+		t.Errorf("Relocation tables are not the same")
+	}
+}
+
+func TestRelocateSegmentsWithHoles(t *testing.T) {
+	segments := memory.NewMemorySegmentManager()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.AddSegment()
+	segments.SegmentSizes = map[uint]uint{0: 3, 2: 3}
+	relocationTable, ok := segments.RelocateSegments()
+
+	if !ok {
+		t.Errorf("Memory segment manager doesn't have segment sizes initialized")
+	}
+
+	expectedTable := []uint{1, 4, 4}
+	if !reflect.DeepEqual(expectedTable, relocationTable) {
+		t.Errorf("Relocation tables are not the same")
+	}
+}
+
+func TestRelocateMemory(t *testing.T) {
+	virtualMachine := vm.NewVirtualMachine()
+	segments := virtualMachine.Segments
+	for i := 0; i < 4; i++ {
+		segments.AddSegment()
+	}
+	segments.Memory.Insert(memory.NewRelocatable(0, 0), memory.NewMaybeRelocatableInt(4613515612218425347))
+	segments.Memory.Insert(memory.NewRelocatable(0, 1), memory.NewMaybeRelocatableInt(5))
+	segments.Memory.Insert(memory.NewRelocatable(0, 2), memory.NewMaybeRelocatableInt(2345108766317314046))
+	segments.Memory.Insert(memory.NewRelocatable(1, 0), memory.NewMaybeRelocatableRelocatable(memory.NewRelocatable(2, 0)))
+	segments.Memory.Insert(memory.NewRelocatable(1, 1), memory.NewMaybeRelocatableRelocatable(memory.NewRelocatable(3, 0)))
+	segments.Memory.Insert(memory.NewRelocatable(1, 5), memory.NewMaybeRelocatableInt(5))
+
+	segments.ComputeEffectiveSizes()
+
+	relocationTable, ok := segments.RelocateSegments()
+	if !ok {
+		t.Errorf("Could not create relocation table")
+	}
+
+	relocatedMemory, err := segments.RelocateMemory(&relocationTable)
+	if err != nil {
+		t.Errorf("Test failed with error: %s", err)
+	}
+
+	expectedMemory := map[uint]uint{1: 4613515612218425347, 2: 5, 3: 2345108766317314046, 4: 10, 5: 10, 9: 5}
+	for i, v := range expectedMemory {
+		actual := relocatedMemory[i]
+		if actual != v {
+			t.Errorf("Expected relocated memory at index %d to be %d but it's %d", i, v, actual)
+		}
+	}
+}

pkg/vm/vm_core.go

@@ -7,5 +7,11 @@ import "github.com/lambdaclass/cairo-vm.go/pkg/vm/memory"
 type VirtualMachine struct {
 	runContext  RunContext
 	currentStep uint
-	segments    memory.MemorySegmentManager
+	Segments    *memory.MemorySegmentManager
+}
+
+func NewVirtualMachine() *VirtualMachine {
+	segments := memory.NewMemorySegmentManager()
+
+	return &VirtualMachine{Segments: segments}
 }