Finalize segments + layout flags (plain, small, starknet_with_keccak) (…

…#432)

* Add finalize method, calculation of builtin segment size, builtins parameters

* Add end_run method and helper methods, modify the builtins values to match layout small

* Fixed usage of constants in builtins, added checking range check limits

* remove vm traces

* Linter errors

* Debug finalization

* Debug range check limits calculation

* Fix builtins test errors

* Debug output segment finalisation, add annotations to what should be moved to the layout

* Modify mock integration test, add annotations to the functions

* Add annotations

* Add new annotations

* Add layouts

* Temporarily modify integration testsuite

* Fix layout problems

* removed unused layout plain

* Layout flag with layouts: plain, small, all_cairo (#435)

* Create layout flag with layouts plain and small

* Fix error handling

* Fix error handling

* Add layout flag to Nethermind vm in the integration tests

* Fix unit tests

* Fix builtins integration tests, add layouts to builtins tests, remove failing "slices" package

* Refactor code, remove unnecessary comments

* Replace unsupported "all_cairo" layout with "starknet_with_keccak"

* Lint the project

* Add builtin metadata and GetAllocatedSize to PoseidonBuiltin

* Add Poseidon Builtin to starknet_with_keccak layout

* Fix integration tests for poseidon builtin

* Address first batch of comments

* use correct constant for poseidon builtin

* Add starknet builtin detection without type assertion

* Add comments for layout struct fields

* Address comments from PR, move calculating allocated size to separate functionality

* Address comment regarding backward iteration on range check usage method

* Add starknet_with_keccak to the integration tests in the python VM run

---------

Co-authored-by: Harikrishnan Shaji <[email protected]>
Co-authored-by: Tristan <[email protected]>
Co-authored-by: cicr99 <[email protected]>

cmd/cli/main.go

@@ -17,7 +17,7 @@ func main() {
 	var entrypointOffset uint64
 	var traceLocation string
 	var memoryLocation string
-
+	var layoutName string
 	app := &cli.App{
 		Name:                 "cairo-vm",
 		Usage:                "A cairo virtual machine",
@@ -61,6 +61,12 @@ func main() {
 						Required:    false,
 						Destination: &memoryLocation,
 					},
+					&cli.StringFlag{
+						Name:        "layout",
+						Usage:       "specifies the set of builtins to be used",
+						Required:    false,
+						Destination: &layoutName,
+					},
 				},
 				Action: func(ctx *cli.Context) error {
 					// TODO: move this action's body to a separate function to decrease the
@@ -89,9 +95,8 @@ func main() {
 					if err != nil {
 						return fmt.Errorf("cannot create hints: %w", err)
 					}
-
 					fmt.Println("Running....")
-					runner, err := runnerzero.NewRunner(program, hints, proofmode, maxsteps)
+					runner, err := runnerzero.NewRunner(program, hints, proofmode, maxsteps, layoutName)
 					if err != nil {
 						return fmt.Errorf("cannot create runner: %w", err)
 					}
@@ -111,6 +116,10 @@ func main() {
 					}
 
 					if proofmode {
+						runner.EndRun()
+						if err := runner.FinalizeSegments(); err != nil {
+							return fmt.Errorf("cannot finalize segments: %w", err)
+						}
 						trace, memory, err := runner.BuildProof()
 						if err != nil {
 							return fmt.Errorf("cannot build proof: %w", err)

integration_tests/builtin_tests/poseidon_test.cairo → integration_tests/builtin_tests/poseidon_test.starknet_with_keccak.cairo

@@ -1,4 +1,4 @@
-%builtins output poseidon range_check
+%builtins output poseidon
 from starkware.cairo.common.cairo_builtins import PoseidonBuiltin
 from starkware.cairo.common.poseidon_state import PoseidonBuiltinState
 
@@ -38,7 +38,7 @@ func test_poseidon_builtin_random_big_values{output_ptr: felt*, poseidon_ptr: Po
     return ();
 }
 
-func main{output_ptr: felt*, poseidon_ptr: PoseidonBuiltin*, range_check_ptr}() {
+func main{output_ptr: felt*, poseidon_ptr: PoseidonBuiltin*}() {
     test_poseidon_builtin_random_small_values();
     test_poseidon_builtin_random_big_values();
     return();

integration_tests/cairo_files/test.starknet_with_keccak.cairo

@@ -0,0 +1,3 @@
+func main() {
+    return ();
+}

integration_tests/cairozero_test.go

@@ -169,6 +169,8 @@ func runPythonVm(testFilename, path string) (string, string, error) {
 	// A file without this suffix will use the default ("plain") layout.
 	if strings.HasSuffix(testFilename, ".small.cairo") {
 		args = append(args, "--layout", "small")
+	} else if strings.HasSuffix(testFilename, ".starknet_with_keccak.cairo") {
+		args = append(args, "--layout", "starknet_with_keccak")
 	}
 
 	cmd := exec.Command("cairo-run", args...)
@@ -189,6 +191,16 @@ func runVm(path string) (string, string, string, error) {
 	traceOutput := swapExtenstion(path, traceSuffix)
 	memoryOutput := swapExtenstion(path, memorySuffix)
 
+	// If any other layouts are needed, add the suffix checks here.
+	// The convention would be: ".$layout.cairo"
+	// A file without this suffix will use the default ("plain") layout, which is a layout with no builtins included"
+	layout := "plain"
+	if strings.Contains(path, ".small") {
+		layout = "small"
+	} else if strings.Contains(path, ".starknet_with_keccak") {
+		layout = "starknet_with_keccak"
+	}
+
 	cmd := exec.Command(
 		"../bin/cairo-vm",
 		"run",
@@ -197,6 +209,8 @@ func runVm(path string) (string, string, string, error) {
 		traceOutput,
 		"--memoryfile",
 		memoryOutput,
+		"--layout",
+		layout,
 		path,
 	)
 
@@ -283,7 +297,7 @@ func memoryRepr(memory []*fp.Element) string {
 }
 
 func TestFailingRangeCheck(t *testing.T) {
-	compiledOutput, err := compileZeroCode("./builtin_tests/range_check.cairo")
+	compiledOutput, err := compileZeroCode("./builtin_tests/range_check.small.cairo")
 	require.NoError(t, err)
 
 	_, _, _, err = runVm(compiledOutput)
@@ -293,7 +307,7 @@ func TestFailingRangeCheck(t *testing.T) {
 }
 
 func TestBitwise(t *testing.T) {
-	compiledOutput, err := compileZeroCode("./builtin_tests/bitwise_builtin_test.cairo")
+	compiledOutput, err := compileZeroCode("./builtin_tests/bitwise_builtin_test.starknet_with_keccak.cairo")
 	require.NoError(t, err)
 
 	_, _, _, err = runVm(compiledOutput)
@@ -303,7 +317,7 @@ func TestBitwise(t *testing.T) {
 }
 
 func TestPedersen(t *testing.T) {
-	compiledOutput, err := compileZeroCode("./builtin_tests/pedersen_test.cairo")
+	compiledOutput, err := compileZeroCode("./builtin_tests/pedersen_test.small.cairo")
 	require.NoError(t, err)
 
 	_, _, output, err := runVm(compiledOutput)
@@ -314,7 +328,7 @@ func TestPedersen(t *testing.T) {
 }
 
 func TestPoseidon(t *testing.T) {
-	compiledOutput, err := compileZeroCode("./builtin_tests/poseidon_test.cairo")
+	compiledOutput, err := compileZeroCode("./builtin_tests/poseidon_test.starknet_with_keccak.cairo")
 	require.NoError(t, err)
 
 	_, _, output, err := runVm(compiledOutput)
@@ -325,7 +339,7 @@ func TestPoseidon(t *testing.T) {
 }
 
 func TestECDSA(t *testing.T) {
-	compiledOutput, err := compileZeroCode("./builtin_tests/ecdsa_test.cairo")
+	compiledOutput, err := compileZeroCode("./builtin_tests/ecdsa_test.starknet_with_keccak.cairo")
 	require.NoError(t, err)
 
 	_, _, _, err = runVm(compiledOutput)
@@ -335,7 +349,7 @@ func TestECDSA(t *testing.T) {
 }
 
 func TestEcOp(t *testing.T) {
-	compiledOutput, err := compileZeroCode("./builtin_tests/ecop.cairo")
+	compiledOutput, err := compileZeroCode("./builtin_tests/ecop.starknet_with_keccak.cairo")
 	require.NoError(t, err)
 
 	_, _, _, err = runVm(compiledOutput)
@@ -346,7 +360,7 @@ func TestEcOp(t *testing.T) {
 }
 
 func TestKeccak(t *testing.T) {
-	compiledOutput, err := compileZeroCode("./builtin_tests/keccak_test.cairo")
+	compiledOutput, err := compileZeroCode("./builtin_tests/keccak_test.starknet_with_keccak.cairo")
 	require.NoError(t, err)
 
 	_, _, output, err := runVm(compiledOutput)

pkg/runners/zero/zero.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/NethermindEth/cairo-vm-go/pkg/hintrunner"
 	"github.com/NethermindEth/cairo-vm-go/pkg/hintrunner/hinter"
+	"github.com/NethermindEth/cairo-vm-go/pkg/parsers/starknet"
 	"github.com/NethermindEth/cairo-vm-go/pkg/utils"
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm"
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm/builtins"
@@ -24,17 +25,22 @@ type ZeroRunner struct {
 	maxsteps  uint64
 	// auxiliar
 	runFinished bool
+	layout      builtins.Layout
 }
 
 // Creates a new Runner of a Cairo Zero program
-func NewRunner(program *Program, hints map[uint64][]hinter.Hinter, proofmode bool, maxsteps uint64) (ZeroRunner, error) {
+func NewRunner(program *Program, hints map[uint64][]hinter.Hinter, proofmode bool, maxsteps uint64, layoutName string) (ZeroRunner, error) {
 	hintrunner := hintrunner.NewHintRunner(hints)
-
+	layout, err := builtins.GetLayout(layoutName)
+	if err != nil {
+		return ZeroRunner{}, err
+	}
 	return ZeroRunner{
 		program:    program,
 		hintrunner: hintrunner,
 		proofmode:  proofmode,
 		maxsteps:   maxsteps,
+		layout:     layout,
 	}, nil
 }
 
@@ -123,7 +129,10 @@ func (runner *ZeroRunner) InitializeMainEntrypoint() (mem.MemoryAddress, error)
 				errors.New("end label not found. Try compiling with `--proof_mode`")
 		}
 
-		stack := runner.initializeBuiltins(memory)
+		stack, err := runner.initializeBuiltins(memory)
+		if err != nil {
+			return mem.UnknownAddress, err
+		}
 		// Add the dummy last fp and pc to the public memory, so that the verifier can enforce [fp - 2] = fp.
 		stack = append([]mem.MemoryValue{mem.MemoryValueFromSegmentAndOffset(
 			vm.ProgramSegment,
@@ -155,7 +164,10 @@ func (runner *ZeroRunner) InitializeMainEntrypoint() (mem.MemoryAddress, error)
 func (runner *ZeroRunner) initializeEntrypoint(
 	initialPCOffset uint64, arguments []*f.Element, returnFp *mem.MemoryValue, memory *mem.Memory,
 ) (mem.MemoryAddress, error) {
-	stack := runner.initializeBuiltins(memory)
+	stack, err := runner.initializeBuiltins(memory)
+	if err != nil {
+		return mem.UnknownAddress, err
+	}
 	for i := range arguments {
 		stack = append(stack, mem.MemoryValueFromFieldElement(arguments[i]))
 	}
@@ -168,14 +180,30 @@ func (runner *ZeroRunner) initializeEntrypoint(
 	}, stack, memory)
 }
 
-func (runner *ZeroRunner) initializeBuiltins(memory *mem.Memory) []mem.MemoryValue {
+func (runner *ZeroRunner) initializeBuiltins(memory *mem.Memory) ([]mem.MemoryValue, error) {
+	builtinsSet := make(map[starknet.Builtin]bool)
+	for _, bRunner := range runner.layout.Builtins {
+		builtinsSet[bRunner.Builtin] = true
+	}
+	// check if all builtins from the program are in the layout
+	for _, programBuiltin := range runner.program.Builtins {
+		if _, found := builtinsSet[programBuiltin]; !found {
+			builtinName, err := programBuiltin.MarshalJSON()
+			if err != nil {
+				return []mem.MemoryValue{}, err
+			}
+			return []mem.MemoryValue{}, fmt.Errorf("builtin %s not found in the layout: %s", builtinName, runner.layout.Name)
+		}
+	}
 	stack := []mem.MemoryValue{}
-	for _, builtin := range runner.program.Builtins {
-		bRunner := builtins.Runner(builtin)
-		builtinSegment := memory.AllocateBuiltinSegment(bRunner)
-		stack = append(stack, mem.MemoryValueFromMemoryAddress(&builtinSegment))
+	// adding to the stack only the builtins that are both in the program and in the layout
+	for _, bRunner := range runner.layout.Builtins {
+		builtinSegment := memory.AllocateBuiltinSegment(bRunner.Runner)
+		if utils.Contains(runner.program.Builtins, bRunner.Builtin) {
+			stack = append(stack, mem.MemoryValueFromMemoryAddress(&builtinSegment))
+		}
 	}
-	return stack
+	return stack, nil
 }
 
 func (runner *ZeroRunner) initializeVm(
@@ -240,6 +268,98 @@ func (runner *ZeroRunner) RunFor(steps uint64) error {
 	return nil
 }
 
+// EndRun is responsible for running the additional steps after the program was executed,
+// until the checkUsedCells doesn't return any 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 *ZeroRunner) EndRun() {
+	if runner.proofmode {
+		for runner.checkUsedCells() != nil {
+			pow2Steps := utils.NextPowerOfTwo(runner.vm.Step + 1)
+			if err := runner.RunFor(pow2Steps); err != nil {
+				panic(err)
+			}
+		}
+	}
+}
+
+// checkUsedCells returns error if not enough steps were made to allocate required number of cells for builtins
+// or there are not enough trace cells to fill the entire range check range
+func (runner *ZeroRunner) checkUsedCells() error {
+	for _, bRunner := range runner.layout.Builtins {
+		builtinSegment, ok := runner.vm.Memory.FindSegmentWithBuiltin(bRunner.Runner.String())
+		if ok {
+			_, err := bRunner.Runner.GetAllocatedSize(builtinSegment.Len(), runner.steps())
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return runner.checkRangeCheckUsage()
+}
+
+func (runner *ZeroRunner) checkRangeCheckUsage() error {
+	rcMin, rcMax := runner.getPermRangeCheckLimits()
+	var rcUnitsUsedByBuiltins uint64
+	for _, builtin := range runner.program.Builtins {
+		if builtin == starknet.RangeCheck {
+			bRunner := builtins.Runner(builtin)
+			rangeCheckRunner, _ := bRunner.(*builtins.RangeCheck)
+			rangeCheckSegment, ok := runner.vm.Memory.FindSegmentWithBuiltin(rangeCheckRunner.String())
+			if ok {
+				rcUnitsUsedByBuiltins += rangeCheckSegment.Len() * rangeCheckRunner.RangeCheckNParts
+			}
+		}
+	}
+	unusedRcUnits := (runner.layout.RcUnits-3)*runner.vm.Step - rcUnitsUsedByBuiltins
+	rcUsageUpperBound := rcMax - rcMin
+	if unusedRcUnits < rcUsageUpperBound {
+		return fmt.Errorf("RangeCheck usage is %d, but the upper bound is %d", unusedRcUnits, rcUsageUpperBound)
+	}
+	return nil
+}
+
+func (runner *ZeroRunner) getPermRangeCheckLimits() (uint64, uint64) {
+	rcMin, rcMax := runner.vm.RcLimitsMin, runner.vm.RcLimitsMax
+
+	for _, builtin := range runner.program.Builtins {
+		if builtin == starknet.RangeCheck {
+			bRunner := builtins.Runner(builtin)
+			rangeCheckRunner, _ := bRunner.(*builtins.RangeCheck)
+			rangeCheckSegment, ok := runner.vm.Memory.FindSegmentWithBuiltin(rangeCheckRunner.String())
+			if ok {
+				rangeCheckUsageMin, rangeCheckUsageMax := rangeCheckRunner.GetRangeCheckUsage(rangeCheckSegment)
+				if rangeCheckUsageMin < rcMin {
+					rcMin = rangeCheckUsageMin
+				}
+				if rangeCheckUsageMax > rcMax {
+					rcMax = rangeCheckUsageMax
+				}
+			}
+		}
+	}
+	return rcMin, rcMax
+}
+
+// FinalizeSegments calculates the final size of the builtins segments,
+// using number of allocated instances and memory cells per builtin instance.
+// Additionally it sets the final size of the program segment to the program size.
+func (runner *ZeroRunner) FinalizeSegments() error {
+	programSize := uint64(len(runner.program.Bytecode))
+	runner.vm.Memory.Segments[vm.ProgramSegment].Finalize(programSize)
+	for _, bRunner := range runner.layout.Builtins {
+		builtinSegment, ok := runner.vm.Memory.FindSegmentWithBuiltin(bRunner.Runner.String())
+		if ok {
+			size, err := bRunner.Runner.GetAllocatedSize(builtinSegment.Len(), runner.vm.Step)
+			if err != nil {
+				return fmt.Errorf("builtin %s: %v", bRunner.Runner.String(), err)
+			}
+			builtinSegment.Finalize(size)
+		}
+	}
+	return nil
+}
+
 func (runner *ZeroRunner) BuildProof() ([]byte, []byte, error) {
 	relocatedTrace, err := runner.vm.ExecutionTrace()
 	if err != nil {

pkg/runners/zero/zero_benchmark_test.go

@@ -233,7 +233,7 @@ func BenchmarkRunnerWithFibonacci(b *testing.B) {
 			panic(err)
 		}
 
-		runner, err := NewRunner(program, hints, true, math.MaxUint64)
+		runner, err := NewRunner(program, hints, true, math.MaxUint64, "plain")
 		if err != nil {
 			panic(err)
 		}