Feature: Proof mode for the VM (#34)

* Implement relocation to get relocated memory

* Implement Proof Mode for VM

* Heal vm tests code

* Finish trace relocation

* remove dead code

* Update gitignore

* Minor improvements overall

* Add test for segment relocation

* minor improvements to trace execution

* bug fix

* Remove rogue prints

* Address review

* Add context relocate function

.gitignore

@@ -1,4 +1,6 @@
 bin/
+.idea/
+.python-version
+
 .DS_Store
 vendor/
-.idea/

pkg/vm/memory/memory_manager.go

@@ -8,14 +8,45 @@ type MemoryManager struct {
 	Memory *Memory
 }
 
-func CreateMemoryManager() (*MemoryManager, error) {
+// Creates a new memory manager
+func CreateMemoryManager() *MemoryManager {
 	memory := InitializeEmptyMemory()
 
 	return &MemoryManager{
 		Memory: memory,
-	}, nil
+	}
 }
 
-func (mm *MemoryManager) GetByteCodeAt(segmentIndex uint64, offset uint64) *f.Element {
-	return nil
+// It returns all segments in memory but relocated as a single segment
+// Each element is a pointer to a field element, if the cell was not accessed,
+// nil is stored instead
+func (mm *MemoryManager) RelocateMemory() []*f.Element {
+	maxMemoryUsed := 0
+	// segmentsOffsets[0] =  0
+	// segmentsOffsets[1] = len(segment[0])
+	// segmentsOffsets[N] = len(segment[n - 1]) + sum of segmentsOffsets[n - i] for i in [0, n-1]
+	segmentsOffsets := make([]uint64, uint64(len(mm.Memory.Segments))+1)
+	for i, segment := range mm.Memory.Segments {
+		maxMemoryUsed += len(segment.Data)
+		segmentsOffsets[i+1] = segmentsOffsets[i] + uint64(len(segment.Data))
+	}
+
+	relocatedMemory := make([]*f.Element, maxMemoryUsed)
+	for i, segment := range mm.Memory.Segments {
+		for j, cell := range segment.Data {
+			var felt *f.Element
+			if !cell.Accessed {
+				continue
+			}
+			if cell.Value.IsAddress() {
+				felt = cell.Value.address.Relocate(segmentsOffsets)
+			} else {
+				felt = cell.Value.felt
+			}
+
+			relocatedMemory[segmentsOffsets[i]+uint64(j)] = felt
+		}
+	}
+
+	return relocatedMemory
 }

pkg/vm/memory/memory_manager_test.go

@@ -0,0 +1,149 @@
+package memory
+
+import (
+	"fmt"
+	"testing"
+
+	f "github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMemoryRelocationWithFelt(t *testing.T) {
+	// segment 0: [2, -, -, 3]
+	// segment 3: [5, -, 7, -, 11, 13]
+	// relocated: [2, -, -, 3, 5, -, 7, -, 11, 13]
+
+	manager := CreateMemoryManager()
+	updateMemoryWithValues(
+		manager.Memory,
+		[]memoryWrite{
+			// segment zero
+			{0, 0, uint64(2)},
+			{0, 3, uint64(3)},
+			// segment three
+			{3, 0, uint64(5)},
+			{3, 2, uint64(7)},
+			{3, 4, uint64(11)},
+			{3, 5, uint64(13)},
+		},
+	)
+
+	res := manager.RelocateMemory()
+
+	expected := []*f.Element{
+		// segment zero
+		new(f.Element).SetUint64(2),
+		nil,
+		nil,
+		new(f.Element).SetUint64(3),
+		// segment three
+		new(f.Element).SetUint64(5),
+		nil,
+		new(f.Element).SetUint64(7),
+		nil,
+		new(f.Element).SetUint64(11),
+		new(f.Element).SetUint64(13),
+	}
+
+	require.Equal(t, len(expected), len(res))
+	require.Equal(t, expected, res)
+}
+
+func TestMemoryRelocationWithAddress(t *testing.T) {
+	// segment 0: [-, 1, -, 1:5] (4)
+	// segment 1: [1, 4:3, 7, -, -, 13] (10)
+	// segment 2: [0:1] (11)
+	// segment 3: [2:0] (12)
+	// segment 4: [0:0, 1:1, 1:5, 15] (16)
+	// relocated: [
+	//      zero:   -,  1, -,  9,
+	//      one:    1, 15, 7,  -, -, 13,
+	//      two:    1,
+	//      three: 10,
+	//      four:   0,  5, 9, 15,
+	// ]
+
+	manager := CreateMemoryManager()
+	updateMemoryWithValues(
+		manager.Memory,
+		[]memoryWrite{
+			// segment zero
+			{0, 1, uint64(1)},
+			{0, 3, NewMemoryAddress(1, 5)},
+			// segment one
+			{1, 0, uint64(1)},
+			{1, 1, NewMemoryAddress(4, 3)},
+			{1, 2, uint64(7)},
+			{1, 5, uint64(13)},
+			// segment two
+			{2, 0, NewMemoryAddress(0, 1)},
+			// segment three
+			{3, 0, NewMemoryAddress(2, 0)},
+			// segment four
+			{4, 0, NewMemoryAddress(0, 0)},
+			{4, 1, NewMemoryAddress(1, 1)},
+			{4, 2, NewMemoryAddress(1, 5)},
+			{4, 3, uint64(15)},
+		},
+	)
+
+	res := manager.RelocateMemory()
+
+	expected := []*f.Element{
+		// segment zero
+		nil,
+		new(f.Element).SetUint64(1),
+		nil,
+		new(f.Element).SetUint64(9),
+		// segment one
+		new(f.Element).SetUint64(1),
+		new(f.Element).SetUint64(15),
+		new(f.Element).SetUint64(7),
+		nil,
+		nil,
+		new(f.Element).SetUint64(13),
+		// segment two
+		new(f.Element).SetUint64(1),
+		// segment three
+		new(f.Element).SetUint64(10),
+		// segment 4
+		new(f.Element).SetUint64(0),
+		new(f.Element).SetUint64(5),
+		new(f.Element).SetUint64(9),
+		new(f.Element).SetUint64(15),
+	}
+
+	require.Equal(t, len(expected), len(res))
+	require.Equal(t, expected, res)
+}
+
+type memoryWrite struct {
+	SegmentIndex uint64
+	Offset       uint64
+	Value        any
+}
+
+func updateMemoryWithValues(memory *Memory, valuesToWrite []memoryWrite) {
+	var max_segment uint64 = 0
+	for _, toWrite := range valuesToWrite {
+		// wrap any inside a memory value
+		val, err := MemoryValueFromAny(toWrite.Value)
+		if err != nil {
+			panic(err)
+		}
+
+		// if the destination segment does not exist, create it
+		for toWrite.SegmentIndex >= max_segment {
+			max_segment += 1
+			memory.AllocateEmptySegment()
+		}
+
+		fmt.Println("c")
+		// write the memory val
+		err = memory.Write(toWrite.SegmentIndex, toWrite.Offset, val)
+		if err != nil {
+			panic(err)
+		}
+
+	}
+}

pkg/vm/memory/memory_value.go

@@ -93,6 +93,13 @@ func (address *MemoryAddress) Sub(lhs *MemoryAddress, rhs any) (*MemoryAddress,
 	}
 }
 
+func (address *MemoryAddress) Relocate(segmentsOffset []uint64) *f.Element {
+	// no risk overflow because this sizes exists in actual Memory
+	// so if by chance the uint64 addition overflowed, then we have
+	// a machine with more than 2**64 bytes of memory (quite a lot!)
+	return new(f.Element).SetUint64(segmentsOffset[address.SegmentIndex] + address.Offset)
+}
+
 func (address MemoryAddress) String() string {
 	return fmt.Sprintf(
 		"Memory Address: segment: %d, offset: %d", address.SegmentIndex, address.Offset,
@@ -137,6 +144,8 @@ func MemoryValueFromSegmentAndOffset[T constraints.Integer](segmentIndex, offset
 
 func MemoryValueFromAny(anyType any) (*MemoryValue, error) {
 	switch t := anyType.(type) {
+	case uint64:
+		return MemoryValueFromInt(anyType.(uint64)), nil
 	case *f.Element:
 		return MemoryValueFromFieldElement(anyType.(*f.Element)), nil
 	case *MemoryAddress:
@@ -279,16 +288,3 @@ func (mv *MemoryValue) Uint64() (uint64, error) {
 
 	return mv.felt.Uint64(), nil
 }
-
-// Note: Commenting this function since relocation is possibly going to look
-// different.
-// Given a map of segment relocation, update a memory address location
-//func (r *MemoryAddress) Relocate(r1 *MemoryAddress, segmentsOffsets *map[uint64]*MemoryAddress) (*MemoryAddress, error) {
-//	if (*segmentsOffsets)[r1.SegmentIndex] == nil {
-//		return nil, fmt.Errorf("missing segment %d relocation rule", r.SegmentIndex)
-//	}
-//
-//	r, err := r.Add((*segmentsOffsets)[r1.SegmentIndex], &MemoryAddress{0, r1.Offset})
-//
-//	return r, err
-//}

pkg/vm/vm.go

@@ -28,48 +28,58 @@ type Context struct {
 	Pc uint64
 }
 
+// relocates pc, ap and fp to be their real address value
+// that is, pc + 0, ap + programSegmentOffset, fp + programSegmentOffset
+func (ctx *Context) Relocate(executionSegmentOffset uint64) {
+	ctx.Ap += executionSegmentOffset
+	ctx.Fp += executionSegmentOffset
+}
+
 // This type represents the current execution context of the vm
 type VirtualMachineConfig struct {
-	Trace bool
-	// Todo(rodro): Update this property to include all builtins
-	Builtins bool
+	// If true, the vm outputs the trace and the relocated memory at the end of execution
+	ProofMode bool
 }
 
 type VirtualMachine struct {
 	Context       Context
 	MemoryManager *mem.MemoryManager
 	Step          uint64
-	Config        VirtualMachineConfig
+	Trace         []Context
+	config        VirtualMachineConfig
 }
 
 // NewVirtualMachine creates a VM from the program bytecode using a specified config.
 func NewVirtualMachine(programBytecode []*f.Element, config VirtualMachineConfig) (*VirtualMachine, error) {
-	manager, err := mem.CreateMemoryManager()
-	if err != nil {
-		return nil, fmt.Errorf("error creating new virtual machine: %w", err)
-	}
-
+	// Initialize memory with to initial segments:
+	// the first one for the program segment and
+	// the second one to keep track of the execution
+	manager := mem.CreateMemoryManager()
 	// 0 (programSegment) <- segment where the bytecode is stored
-	_, err = manager.Memory.AllocateSegment(programBytecode)
+	_, err := manager.Memory.AllocateSegment(programBytecode)
 	if err != nil {
 		return nil, fmt.Errorf("error loading bytecode: %w", err)
 	}
-
 	// 1 (executionSegment) <- segment where ap and fp move around
-	manager.Memory.Segments = append(manager.Memory.Segments, mem.EmptySegmentWithCapacity(10))
+	manager.Memory.AllocateEmptySegment()
+
+	// Initialize the trace if necesary
+	var trace []Context
+	if config.ProofMode {
+		trace = make([]Context, 0)
+	}
 
 	return &VirtualMachine{
 		Context:       Context{Fp: 0, Ap: 0, Pc: 0},
 		Step:          0,
 		MemoryManager: manager,
-		Config:        config,
+		Trace:         trace,
+		config:        config,
 	}, nil
 }
 
 // todo(rodro): add a cache mechanism for not decoding the same instruction twice
 
-// todo(rodro): how to know when te execute a hint or normal instruction
-
 func (vm *VirtualMachine) RunStep(hintRunner HintRunner) error {
 	// Run hint
 	err := hintRunner.RunHint(vm)
@@ -93,6 +103,11 @@ func (vm *VirtualMachine) RunStep(hintRunner HintRunner) error {
 		return fmt.Errorf("cannot decode step at %d: %w", vm.Context.Pc, err)
 	}
 
+	// store the trace before state change
+	if vm.config.ProofMode {
+		vm.Trace = append(vm.Trace, vm.Context)
+	}
+
 	err = vm.RunInstruction(instruction)
 	if err != nil {
 		return fmt.Errorf("cannot run step at %d: %w", vm.Context.Pc, err)
@@ -107,8 +122,6 @@ func (vm *VirtualMachine) RunStepAt(hinter HintRunner, pc uint64) error {
 }
 
 func (vm *VirtualMachine) RunInstruction(instruction *Instruction) error {
-	// todo(rodro): any OffOpX can be negative, a better math system is required due to
-	// substraction. Also it will need to handle overflows and underflows
 	dstCell, err := vm.getCellDst(instruction)
 	if err != nil {
 		return err
@@ -162,6 +175,23 @@ func (vm *VirtualMachine) RunInstruction(instruction *Instruction) error {
 	return nil
 }
 
+// It returns the current trace entry, the public memory, and the occurrence of an error
+func (vm *VirtualMachine) Proof() ([]Context, []*f.Element, error) {
+	if !vm.config.ProofMode {
+		return nil, nil, fmt.Errorf("cannot get proof if proof mode is off")
+	}
+
+	totalBytecode := vm.MemoryManager.Memory.Segments[ProgramSegment].Len()
+	for i := range vm.Trace {
+		vm.Trace[i].Relocate(totalBytecode)
+	}
+
+	// after that, get the relocated memory
+	relocatedMemory := vm.MemoryManager.RelocateMemory()
+
+	return vm.Trace, relocatedMemory, nil
+}
+
 func (vm *VirtualMachine) getCellDst(instruction *Instruction) (*mem.Cell, error) {
 	var dstRegister uint64
 	if instruction.DstRegister == Ap {