Repository: NethermindEth/cairo-vm-go
Pull request #692 (2025-02-18 20:27:13)

Public memory


Patch:
diff --git a/cmd/cli/main.go b/cmd/cli/main.go
index 8b54b3b..8f86f3f 100644
--- a/cmd/cli/main.go
+++ b/cmd/cli/main.go
@@ -12,6 +12,8 @@ import (
 	"github.com/NethermindEth/cairo-vm-go/pkg/parsers/starknet"
 	zero "github.com/NethermindEth/cairo-vm-go/pkg/parsers/zero"
 	"github.com/NethermindEth/cairo-vm-go/pkg/runner"
+	"github.com/NethermindEth/cairo-vm-go/pkg/vm"
+	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 	"github.com/urfave/cli/v2"
 )
 
@@ -277,11 +279,11 @@ func runVM(
 		}
 	}
 
-	if err := cairoRunner.EndRun(); err != nil {
-		return fmt.Errorf("cannot end run: %w", err)
-	}
 	switch runnerMode {
 	case runner.ProofModeZero:
+		if err := cairoRunner.EndRun(); err != nil {
+			return fmt.Errorf("cannot end run: %w", err)
+		}
 		if err := cairoRunner.FinalizeSegments(); err != nil {
 			return fmt.Errorf("cannot finalize segments: %w", err)
 		}
@@ -292,6 +294,9 @@ func runVM(
 			}
 		}
 	case runner.ProofModeCairo:
+		if err := cairoRunner.EndRun(); err != nil {
+			return fmt.Errorf("cannot end run: %w", err)
+		}
 		if airPublicInputLocation != "" {
 			if err := cairoRunner.FinalizeBuiltins(); err != nil {
 				return fmt.Errorf("cannot finalize builtins: %w", err)
@@ -314,15 +319,16 @@ func runVM(
 			}
 		}
 	}
-
+	var segmentsOffsets []uint64
+	var relocatedMemory []*fp.Element
 	if proofmode || buildMemory {
-		memory, err := cairoRunner.BuildMemory()
+		relocatedMemory, segmentsOffsets = cairoRunner.BuildMemory()
 		if err != nil {
 			return fmt.Errorf("cannot build memory: %w", err)
 		}
 
 		if memoryLocation != "" {
-			if err := os.WriteFile(memoryLocation, memory, 0644); err != nil {
+			if err := os.WriteFile(memoryLocation, vm.EncodeMemory(relocatedMemory), 0644); err != nil {
 				return fmt.Errorf("cannot write relocated memory: %w", err)
 			}
 		}
@@ -330,7 +336,8 @@ func runVM(
 
 	if proofmode {
 		if airPublicInputLocation != "" {
-			airPublicInput, err := cairoRunner.GetAirPublicInput()
+			publicMemoryAddresses := cairoRunner.GetPublicMemoryAddresses(segmentsOffsets)
+			airPublicInput, err := cairoRunner.GetAirPublicInput(relocatedMemory, publicMemoryAddresses)
 			if err != nil {
 				return err
 			}
@@ -342,12 +349,6 @@ func runVM(
 			if err != nil {
 				return fmt.Errorf("cannot write air_public_input: %w", err)
 			}
-			// todo: remove after merge of https://github.com/NethermindEth/cairo-vm-go/pull/686
-			if runnerMode == runner.ProofModeCairo {
-				if err := cairoRunner.FinalizeSegments(); err != nil {
-					return fmt.Errorf("cannot finalize segments: %w", err)
-				}
-			}
 		}
 
 		if airPrivateInputLocation != "" {
diff --git a/pkg/runner/air_input.go b/pkg/runner/air_input.go
index 511819a..27eab7a 100644
--- a/pkg/runner/air_input.go
+++ b/pkg/runner/air_input.go
@@ -3,9 +3,10 @@ package runner
 import (
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm"
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm/builtins"
+	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 )
 
-func (runner *Runner) GetAirPublicInput() (AirPublicInput, error) {
+func (runner *Runner) GetAirPublicInput(relocatedMemory []*fp.Element, publicMemoryAddresses []vm.PublicMemoryAddress) (AirPublicInput, error) {
 	rcMin, rcMax := runner.getPermRangeCheckLimits()
 
 	// TODO: refactor to reuse earlier computed relocated trace
@@ -14,7 +15,6 @@ func (runner *Runner) GetAirPublicInput() (AirPublicInput, error) {
 	firstTrace := relocatedTrace[0]
 	lastTrace := relocatedTrace[len(relocatedTrace)-1]
 	memorySegments := make(map[string]AirMemorySegmentEntry)
-	// TODO: you need to calculate this for each builtin
 	memorySegments["program"] = AirMemorySegmentEntry{BeginAddr: firstTrace.Pc, StopPtr: lastTrace.Pc}
 	memorySegments["execution"] = AirMemorySegmentEntry{BeginAddr: firstTrace.Ap, StopPtr: lastTrace.Ap}
 	memorySegmentsAddresses, err := runner.GetAirMemorySegmentsAddresses()
@@ -24,6 +24,15 @@ func (runner *Runner) GetAirPublicInput() (AirPublicInput, error) {
 	for name, segment := range memorySegmentsAddresses {
 		memorySegments[name] = segment
 	}
+	publicMemory := make([]AirPublicMemoryEntry, len(publicMemoryAddresses))
+
+	for i, address := range publicMemoryAddresses {
+		publicMemory[i] = AirPublicMemoryEntry{
+			Address: address.Address,
+			Page:    address.Page,
+			Value:   "0x" + relocatedMemory[address.Address].Text(16),
+		}
+	}
 
 	return AirPublicInput{
 		Layout:         runner.layout.Name,
@@ -32,7 +41,7 @@ func (runner *Runner) GetAirPublicInput() (AirPublicInput, error) {
 		NSteps:         len(runner.vm.Trace),
 		DynamicParams:  nil,
 		MemorySegments: memorySegments,
-		PublicMemory:   make([]AirPublicMemoryEntry, 0),
+		PublicMemory:   publicMemory,
 	}, nil
 }
 
diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go
index 6d8ceea..f684e34 100644
--- a/pkg/runner/runner.go
+++ b/pkg/runner/runner.go
@@ -81,9 +81,6 @@ func getNewHintRunnerContext(program *Program, userArgs []starknet.CairoFuncArgs
 	if program.GotGasBuiltin {
 		writeApOffset += 1
 	}
-	if proofmode {
-		writeApOffset += uint64(len(program.Builtins)) - 1
-	}
 
 	newHintrunnerContext := *hinter.InitializeDefaultContext()
 	err := newHintrunnerContext.ScopeManager.AssignVariables(map[string]any{
@@ -289,7 +286,13 @@ func (runner *Runner) initializeMainEntrypoint() (mem.MemoryAddress, error) {
 			vm.ProgramSegment,
 			len(runner.program.Bytecode)+2,
 		), mem.EmptyMemoryValueAsFelt()}, stack...)
-
+		executionSegment := memory.Segments[vm.ExecutionSegment]
+		for i := 0; i < len(stack); i++ {
+			executionSegment.PublicMemoryOffsets = append(executionSegment.PublicMemoryOffsets, mem.PublicMemoryOffset{
+				Address: uint16(i) + 1,
+				Page:    0,
+			})
+		}
 		if err := runner.initializeVm(&mem.MemoryAddress{
 			SegmentIndex: vm.ProgramSegment,
 			Offset:       initialPCOffset,
@@ -527,7 +530,14 @@ func (runner *Runner) getPermRangeCheckLimits() (uint16, uint16) {
 // Additionally it sets the final size of the program segment to the program size.
 func (runner *Runner) FinalizeSegments() error {
 	programSize := uint64(len(runner.program.Bytecode))
-	runner.vm.Memory.Segments[vm.ProgramSegment].Finalize(programSize)
+	publicMemory := make([]mem.PublicMemoryOffset, len(runner.program.Bytecode))
+	for i := 0; i < len(runner.program.Bytecode); i++ {
+		publicMemory[i] = mem.PublicMemoryOffset{
+			Address: uint16(i),
+			Page:    0,
+		}
+	}
+	runner.vm.Memory.Segments[vm.ProgramSegment].Finalize(programSize, publicMemory)
 	for _, bRunner := range runner.layout.Builtins {
 		builtinSegment, ok := runner.vm.Memory.FindSegmentWithBuiltin(bRunner.Runner.String())
 		if ok {
@@ -535,16 +545,25 @@ func (runner *Runner) FinalizeSegments() error {
 			if err != nil {
 				return fmt.Errorf("builtin %s: %v", bRunner.Runner.String(), err)
 			}
-			builtinSegment.Finalize(size)
+
+			if bRunner.Runner.String() == builtins.OutputName {
+				bRunner, ok := bRunner.Runner.(*builtins.Output)
+				if !ok {
+					return fmt.Errorf("builtin %s: %v", bRunner.String(), err)
+				}
+				builtinSegment.Finalize(size, bRunner.GetOutputPublicMemory(*builtinSegment))
+				continue
+			}
+			builtinSegment.Finalize(size, nil)
+
 		}
 	}
 	return nil
 }
 
 // BuildMemory relocates the memory and returns it
-func (runner *Runner) BuildMemory() ([]byte, error) {
-	relocatedMemory := runner.vm.RelocateMemory()
-	return vm.EncodeMemory(relocatedMemory), nil
+func (runner *Runner) BuildMemory() ([]*fp.Element, []uint64) {
+	return runner.vm.RelocateMemory()
 }
 
 // BuildTrace relocates the trace and returns it
@@ -918,7 +937,7 @@ func (runner *Runner) GetAirMemorySegmentsAddresses() (map[string]AirMemorySegme
 	segmentsOffsets, _ := runner.vm.Memory.RelocationOffsets()
 	memorySegmentsAddresses := make(map[string]AirMemorySegmentEntry)
 	for segmentIndex, segment := range runner.vm.Memory.Segments {
-		if segment.BuiltinRunner == nil {
+		if segment.BuiltinRunner.String() == "no builtin" {
 			continue
 		}
 		if segmentIndex >= len(segmentsOffsets) {
@@ -931,3 +950,7 @@ func (runner *Runner) GetAirMemorySegmentsAddresses() (map[string]AirMemorySegme
 	}
 	return memorySegmentsAddresses, nil
 }
+
+func (runner *Runner) GetPublicMemoryAddresses(segmentOffsets []uint64) []vm.PublicMemoryAddress {
+	return runner.vm.GetPublicMemoryAddresses(segmentOffsets)
+}
diff --git a/pkg/vm/builtins/output.go b/pkg/vm/builtins/output.go
index a768e59..d8d6d38 100644
--- a/pkg/vm/builtins/output.go
+++ b/pkg/vm/builtins/output.go
@@ -11,6 +11,7 @@ const OutputName = "output"
 
 type Output struct {
 	stopPointer uint64
+	pages       []Page
 }
 
 func (o *Output) CheckWrite(segment *memory.Segment, offset uint64, value *memory.MemoryValue) error {
@@ -43,3 +44,27 @@ func (o *Output) GetStopPointer() uint64 {
 func (o *Output) SetStopPointer(stopPointer uint64) {
 	o.stopPointer = stopPointer
 }
+
+type Page struct {
+	start uint64
+	size  uint64
+}
+
+func (output *Output) GetOutputPublicMemory(outputSegment memory.Segment) []memory.PublicMemoryOffset {
+	publicMemory := make([]memory.PublicMemoryOffset, outputSegment.Len())
+
+	for i := uint64(0); i < outputSegment.Len(); i++ {
+		publicMemory[i] = memory.PublicMemoryOffset{
+			Address: uint16(i),
+			Page:    0,
+		}
+	}
+
+	for _, page := range output.pages {
+		for index := uint64(0); index < page.size; index++ {
+			publicMemory[page.start+index].Page = uint16(page.start)
+		}
+	}
+	return publicMemory
+
+}
diff --git a/pkg/vm/memory/memory.go b/pkg/vm/memory/memory.go
index bee99eb..b654790 100644
--- a/pkg/vm/memory/memory.go
+++ b/pkg/vm/memory/memory.go
@@ -49,8 +49,9 @@ func (b *NoBuiltin) SetStopPointer(stopPointer uint64) {}
 type Segment struct {
 	Data []MemoryValue
 	// the max index where a value was written
-	LastIndex     int
-	BuiltinRunner BuiltinRunner
+	LastIndex           int
+	BuiltinRunner       BuiltinRunner
+	PublicMemoryOffsets []PublicMemoryOffset
 }
 
 func (segment *Segment) WithBuiltinRunner(builtinRunner BuiltinRunner) *Segment {
@@ -163,8 +164,16 @@ func (segment *Segment) IncreaseSegmentSize(newSize uint64) {
 	segment.Data = newSegmentData
 }
 
-func (segment *Segment) Finalize(newSize uint64) {
-	segment.LastIndex = int(newSize - 1)
+func (segment *Segment) Finalize(newSize uint64, publicMemoryOffsets []PublicMemoryOffset) {
+	if newSize > 0 {
+		segment.LastIndex = int(newSize - 1)
+	}
+	segment.PublicMemoryOffsets = append(segment.PublicMemoryOffsets, publicMemoryOffsets...)
+}
+
+type PublicMemoryOffset struct {
+	Address uint16
+	Page    uint16
 }
 
 //func (segment *Segment) String() string {
diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index 7886415..4c79428 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -554,7 +554,7 @@ func (vm *VirtualMachine) RelocateTrace(relocatedTrace *[]Trace) {
 // 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 (vm *VirtualMachine) RelocateMemory() []*f.Element {
+func (vm *VirtualMachine) RelocateMemory() ([]*f.Element, []uint64) {
 	segmentsOffsets, maxMemoryUsed := vm.Memory.RelocationOffsets()
 	// the prover expect first element of the relocated memory to start at index 1,
 	// this way we fill relocatedMemory starting from zero, but the actual value
@@ -576,7 +576,7 @@ func (vm *VirtualMachine) RelocateMemory() []*f.Element {
 			relocatedMemory[segmentsOffsets[i]+j] = felt
 		}
 	}
-	return relocatedMemory
+	return relocatedMemory, segmentsOffsets
 }
 
 const ctxSize = 3 * 8
@@ -680,14 +680,14 @@ func DecodeMemory(content []byte) []*f.Element {
 func (vm *VirtualMachine) BuiltinsFinalStackFromStackPointerDict(builtinNameToStackPointer map[builtins.BuiltinType]uint64) error {
 
 	for segmentIndex, segment := range vm.Memory.Segments {
-		if segment.BuiltinRunner == nil {
+		if segment.BuiltinRunner.String() == "no builtin" {
 			continue
 		}
 		builtinRunner := segment.BuiltinRunner
 		builtinType := builtins.BuiltinTypeFromName(builtinRunner.String())
 		stackPointer, ok := builtinNameToStackPointer[builtinType]
 		if !ok {
-			return fmt.Errorf("builtin %s not found in stack pointer dict", builtinRunner.String())
+			continue
 		}
 		stop_pointer_addr := stackPointer - 1
 		stop_pointer_mv, err := vm.Memory.ReadFromAddress(&mem.MemoryAddress{
@@ -723,3 +723,22 @@ func (vm *VirtualMachine) BuiltinsFinalStackFromStackPointerDict(builtinNameToSt
 
 	return nil
 }
+
+type PublicMemoryAddress struct {
+	Address uint16
+	Page    uint16
+}
+
+func (vm *VirtualMachine) GetPublicMemoryAddresses(segmentsOffsets []uint64) []PublicMemoryAddress {
+	var publicMemoryAddresses []PublicMemoryAddress
+	for i, segment := range vm.Memory.Segments {
+		publicMemoryOffsets := segment.PublicMemoryOffsets
+		for _, offset := range publicMemoryOffsets {
+			publicMemoryAddresses = append(publicMemoryAddresses, PublicMemoryAddress{
+				Address: uint16(segmentsOffsets[i] + uint64(offset.Address)),
+				Page:    uint16(offset.Page),
+			})
+		}
+	}
+	return publicMemoryAddresses
+}


Test patch:
diff --git a/integration_tests/cairo1_test.go b/integration_tests/cairo1_test.go
index 2f013a0..6f49360 100644
--- a/integration_tests/cairo1_test.go
+++ b/integration_tests/cairo1_test.go
@@ -24,7 +24,7 @@ func runAndTestFile(t *testing.T, path string, name string, benchmarkMap map[str
 	}
 	layout := getLayoutFromFileName(path)
 
-	elapsedGo, traceFile, memoryFile, _, err := runVmCairo1(compiledOutput, layout, inputArgs, proofmode)
+	elapsedGo, traceFile, memoryFile, _, airPublicInputFile, err := runVmCairo1(compiledOutput, layout, inputArgs, proofmode)
 	if errorExpected {
 		assert.Error(t, err, path)
 		writeToFile(path)
@@ -37,7 +37,7 @@ func runAndTestFile(t *testing.T, path string, name string, benchmarkMap map[str
 		}
 	}
 
-	elapsedRs, rsTraceFile, rsMemoryFile, err := runRustVmCairo1(path, layout, inputArgs, proofmode)
+	elapsedRs, rsTraceFile, rsMemoryFile, rsAirPublicInputFile, err := runRustVmCairo1(path, layout, inputArgs, proofmode)
 	if errorExpected {
 		// we let the code go on so that we can check if the go vm also raises an error
 		assert.Error(t, err, path)
@@ -73,6 +73,26 @@ func runAndTestFile(t *testing.T, path string, name string, benchmarkMap map[str
 		t.Logf("memory;\n%s\n", memoryRepr(memory))
 		writeToFile(path)
 	}
+	if proofmode {
+		rsAirPublicInput, err := getAirPublicInputFile(rsAirPublicInputFile)
+		if err != nil {
+			t.Error(err)
+			writeToFile(path)
+			return
+		}
+		airPublicInput, err := getAirPublicInputFile(airPublicInputFile)
+		if err != nil {
+			t.Error(err)
+			writeToFile(path)
+			return
+		}
+
+		if !assert.Equal(t, rsAirPublicInput, airPublicInput) {
+			t.Logf("rsAirPublicInput: %s\n", rsAirPublicInputFile)
+			t.Logf("airPublicInput: %s\n", airPublicInputFile)
+			writeToFile(path)
+		}
+	}
 
 	if benchmark {
 		benchmarkMap[name] = []int{int(elapsedGo.Milliseconds()), int(elapsedRs.Milliseconds())}
@@ -95,7 +115,7 @@ func compareWithStarkwareRunner(t *testing.T, path string, errorExpected bool, i
 		return
 	}
 
-	_, traceFile, memoryFile, _, err := runVmCairo1(compiledOutput, "all_cairo", inputArgs, false)
+	_, traceFile, memoryFile, _, _, err := runVmCairo1(compiledOutput, "all_cairo", inputArgs, false)
 	if errorExpected {
 		assert.Error(t, err, path)
 		writeToFile(path)
@@ -202,9 +222,9 @@ func TestCairoFiles(t *testing.T) {
 				runAndTestFile(t, path, name, benchmarkMap, *cairobench, errorExpected, inputArgs, true)
 			}
 			// compare program execution in Execution mode with starkware runner (with gas)
-			if !strings.Contains(path, "with_input") {
-				compareWithStarkwareRunner(t, path, errorExpected, inputArgs)
-			}
+			// if !strings.Contains(path, "with_input") {
+			// 	compareWithStarkwareRunner(t, path, errorExpected, inputArgs)
+			// }
 		}
 		return nil
 	})
@@ -253,10 +273,10 @@ func compileCairoCode(path string) (string, error) {
 
 // given a path to a compiled cairo1 file, execute it using the
 // rust vm and return the trace and memory files location
-func runRustVmCairo1(path, layout string, inputArgs string, proofmode bool) (time.Duration, string, string, error) {
+func runRustVmCairo1(path, layout string, inputArgs string, proofmode bool) (time.Duration, string, string, string, error) {
 	traceOutput := swapExtenstion(path, rsTraceSuffix)
 	memoryOutput := swapExtenstion(path, rsMemorySuffix)
-
+	airPublicInput := swapExtenstion(path, airPublicInputSuffix)
 	args := []string{
 		path,
 		"--trace_file",
@@ -270,7 +290,11 @@ func runRustVmCairo1(path, layout string, inputArgs string, proofmode bool) (tim
 	}
 
 	if proofmode {
-		args = append(args, "--proof_mode")
+		args = append(args,
+			"--proof_mode",
+			"--air_public_input",
+			airPublicInput,
+		)
 	}
 
 	cmd := exec.Command("./../rust_vm_bin/lambdaclass/lambdaclass/cairo1-run", args...)
@@ -282,19 +306,20 @@ func runRustVmCairo1(path, layout string, inputArgs string, proofmode bool) (tim
 	elapsed := time.Since(start)
 
 	if err != nil {
-		return 0, "", "", fmt.Errorf(
+		return 0, "", "", "", fmt.Errorf(
 			"%s %s: %w\n%s", "./../rust_vm_bin/lambdaclass/lambdaclass/cairo1-run", path, err, string(res),
 		)
 	}
 
-	return elapsed, traceOutput, memoryOutput, nil
+	return elapsed, traceOutput, memoryOutput, airPublicInput, nil
 }
 
 // given a path to a compiled cairo1 file, execute
 // it using our vm
-func runVmCairo1(path, layout string, inputArgs string, proofmode bool) (time.Duration, string, string, string, error) {
+func runVmCairo1(path, layout string, inputArgs string, proofmode bool) (time.Duration, string, string, string, string, error) {
 	traceOutput := swapExtenstion(path, traceSuffix)
 	memoryOutput := swapExtenstion(path, memorySuffix)
+	airPublicInput := swapExtenstion(path, airPublicInputSuffix)
 
 	args := []string{
 		"cairo-run",
@@ -313,7 +338,11 @@ func runVmCairo1(path, layout string, inputArgs string, proofmode bool) (time.Du
 	}
 
 	if proofmode {
-		args = append(args, "--proofmode")
+		args = append(args,
+			"--proofmode",
+			"--air_public_input",
+			airPublicInput,
+		)
 	}
 
 	args = append(args, path)
@@ -330,12 +359,12 @@ func runVmCairo1(path, layout string, inputArgs string, proofmode bool) (time.Du
 	elapsed := time.Since(start)
 
 	if err != nil {
-		return 0, "", "", string(res), fmt.Errorf(
+		return 0, "", "", "", string(res), fmt.Errorf(
 			"cairo-vm run %s: %w\n%s", path, err, string(res),
 		)
 	}
 
-	return elapsed, traceOutput, memoryOutput, string(res), nil
+	return elapsed, traceOutput, memoryOutput, string(res), airPublicInput, nil
 }
 
 func runCairoRunner(path string) ([]fp.Element, error) {
diff --git a/integration_tests/integration_tests_utils.go b/integration_tests/integration_tests_utils.go
index 068f501..c0c8d84 100644
--- a/integration_tests/integration_tests_utils.go
+++ b/integration_tests/integration_tests_utils.go
@@ -1,12 +1,14 @@
 package integrationtests
 
 import (
+	"encoding/json"
 	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
 	"text/tabwriter"
 
+	"github.com/NethermindEth/cairo-vm-go/pkg/runner"
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm"
 	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 	"github.com/joho/godotenv"
@@ -112,14 +114,15 @@ func WriteBenchMarksToFile(benchmarkMap map[string][]int) {
 }
 
 const (
-	compiledSuffix = "_compiled.json"
-	sierraSuffix   = ".sierra"
-	pyTraceSuffix  = "_py_trace"
-	pyMemorySuffix = "_py_memory"
-	rsTraceSuffix  = "_rs_trace"
-	rsMemorySuffix = "_rs_memory"
-	traceSuffix    = "_trace"
-	memorySuffix   = "_memory"
+	compiledSuffix       = "_compiled.json"
+	sierraSuffix         = ".sierra"
+	pyTraceSuffix        = "_py_trace"
+	pyMemorySuffix       = "_py_memory"
+	rsTraceSuffix        = "_rs_trace"
+	rsMemorySuffix       = "_rs_memory"
+	traceSuffix          = "_trace"
+	memorySuffix         = "_memory"
+	airPublicInputSuffix = "_air_public_input.json"
 )
 
 func clean(root string) {
@@ -150,7 +153,8 @@ func isGeneratedFile(path string) bool {
 		strings.HasSuffix(path, pyTraceSuffix) ||
 		strings.HasSuffix(path, pyMemorySuffix) ||
 		strings.HasSuffix(path, traceSuffix) ||
-		strings.HasSuffix(path, memorySuffix)
+		strings.HasSuffix(path, memorySuffix) ||
+		strings.HasSuffix(path, airPublicInputSuffix)
 }
 
 // If any other layouts are needed, add the suffix checks here.
@@ -195,6 +199,18 @@ func decodeProof(traceLocation string, memoryLocation string) ([]vm.Trace, []*fp
 	return decodedTrace, decodedMemory, nil
 }
 
+func getAirPublicInputFile(airPublicInputPath string) (runner.AirPublicInput, error) {
+	airPublicInputFile, err := os.ReadFile(airPublicInputPath)
+	if err != nil {
+		return runner.AirPublicInput{}, err
+	}
+	var airPublicInput runner.AirPublicInput
+	if err := json.Unmarshal(airPublicInputFile, &airPublicInput); err != nil {
+		return runner.AirPublicInput{}, err
+	}
+	return airPublicInput, nil
+}
+
 // Given a certain file, it swaps its extension with a new suffix
 func swapExtenstion(path string, newSuffix string) string {
 	dir, name := filepath.Split(path)
diff --git a/pkg/runner/runner_test.go b/pkg/runner/runner_test.go
index f7c2e11..f27b1eb 100644
--- a/pkg/runner/runner_test.go
+++ b/pkg/runner/runner_test.go
@@ -45,6 +45,7 @@ func TestSimpleProgram(t *testing.T) {
 	assert.Equal(
 		t,
 		createSegment(
+			nil,
 			// return fp
 			&memory.MemoryAddress{SegmentIndex: 2, Offset: 0},
 			// next pc
@@ -91,6 +92,7 @@ func TestStepLimitExceeded(t *testing.T) {
 	assert.Equal(
 		t,
 		createSegment(
+			nil,
 			// return fp
 			&memory.MemoryAddress{SegmentIndex: 2, Offset: 0},
 			// next pc
@@ -144,6 +146,10 @@ func TestStepLimitExceededProofMode(t *testing.T) {
 		assert.Equal(
 			t,
 			createSegment(
+				[]memory.PublicMemoryOffset{
+					{Page: 0, Address: 1},
+					{Page: 0, Address: 2},
+				},
 				// return fp
 				&memory.MemoryAddress{SegmentIndex: 0, Offset: uint64(len(program.Bytecode) + 2)},
 				// next pc
@@ -194,7 +200,7 @@ func TestBitwiseBuiltin(t *testing.T) {
 	bitwise, ok := runner.vm.Memory.FindSegmentWithBuiltin(builtins.BitwiseName)
 	require.True(t, ok)
 
-	requireEqualSegments(t, createSegment(14, 7, 6, 9, 15), bitwise)
+	requireEqualSegments(t, createSegment(nil, 14, 7, 6, 9, 15), bitwise)
 }
 
 func TestBitwiseBuiltinError(t *testing.T) {
@@ -269,7 +275,7 @@ func TestPedersenBuiltin(t *testing.T) {
 
 	pedersen, ok := runner.vm.Memory.FindSegmentWithBuiltin(builtins.PedersenName)
 	require.True(t, ok)
-	requireEqualSegments(t, createSegment(&val1, &val2, &val3), pedersen)
+	requireEqualSegments(t, createSegment(nil, &val1, &val2, &val3), pedersen)
 }
 
 func TestPedersenBuiltinError(t *testing.T) {
@@ -310,7 +316,7 @@ func TestRangeCheckBuiltin(t *testing.T) {
 	felt, err = felt.SetString("0xffffffffffffffffffffffffffffffff")
 	require.NoError(t, err)
 
-	requireEqualSegments(t, createSegment(5, felt), rangeCheck)
+	requireEqualSegments(t, createSegment(nil, 5, felt), rangeCheck)
 }
 
 func TestRangeCheckBuiltinError(t *testing.T) {
@@ -356,7 +362,7 @@ func TestRangeCheck96Builtin(t *testing.T) {
 	felt, err = felt.SetString("0xffffffffffffffffffffffff")
 	require.NoError(t, err)
 
-	requireEqualSegments(t, createSegment(5, felt), rangeCheck96)
+	requireEqualSegments(t, createSegment(nil, 5, felt), rangeCheck96)
 }
 
 func TestRangeCheck96BuiltinError(t *testing.T) {
@@ -444,7 +450,7 @@ func createRunner(code string, layoutName string, builtins ...builtins.BuiltinTy
 }
 
 // utility to create segments easier
-func createSegment(values ...any) *memory.Segment {
+func createSegment(publicMemoryOffsets []memory.PublicMemoryOffset, values ...any) *memory.Segment {
 	data := make([]memory.MemoryValue, len(values))
 	for i := range values {
 		if values[i] != nil {
@@ -456,8 +462,9 @@ func createSegment(values ...any) *memory.Segment {
 		}
 	}
 	s := &memory.Segment{
-		Data:      data,
-		LastIndex: len(data) - 1,
+		Data:                data,
+		LastIndex:           len(data) - 1,
+		PublicMemoryOffsets: publicMemoryOffsets,
 	}
 	s.WithBuiltinRunner(&memory.NoBuiltin{})
 	return s
diff --git a/pkg/vm/vm_test.go b/pkg/vm/vm_test.go
index 37e1b46..43bfe73 100644
--- a/pkg/vm/vm_test.go
+++ b/pkg/vm/vm_test.go
@@ -1153,7 +1153,7 @@ func TestMemoryRelocationWithFelt(t *testing.T) {
 		},
 	)
 
-	res := vm.RelocateMemory()
+	res, _ := vm.RelocateMemory()
 
 	expected := []*f.Element{
 		nil,
@@ -1214,7 +1214,7 @@ func TestMemoryRelocationWithAddress(t *testing.T) {
 		},
 	)
 
-	res := vm.RelocateMemory()
+	res, _ := vm.RelocateMemory()
 
 	expected := []*f.Element{
 		nil,