Tests: Add unit tests for Cairo Zero using Builtins (#125)

* Update starknet parser pkg name

* Add builtin tests

* Overall minor refactoring

pkg/parsers/starknet/hint.go

@@ -1,4 +1,4 @@
-package vm
+package starknet
 
 import (
 	"encoding/json"

pkg/parsers/starknet/hint_test.go

@@ -1,4 +1,4 @@
-package vm
+package starknet
 
 import (
 	"encoding/json"

pkg/parsers/starknet/starknet.go

@@ -1,4 +1,4 @@
-package vm
+package starknet
 
 import (
 	"encoding/json"

pkg/parsers/starknet/starknet_test.go

@@ -1,4 +1,4 @@
-package vm
+package starknet
 
 import (
 	"testing"

pkg/runners/zero/program.go

@@ -4,7 +4,7 @@ import (
 	"errors"
 	"fmt"
 
-	starknetParser "github.com/NethermindEth/cairo-vm-go/pkg/parsers/starknet"
+	sn "github.com/NethermindEth/cairo-vm-go/pkg/parsers/starknet"
 	"github.com/NethermindEth/cairo-vm-go/pkg/parsers/zero"
 	f "github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
 )
@@ -17,7 +17,7 @@ type Program struct {
 	// it stores the start and end label pcs
 	Labels map[string]uint64
 	// builtins
-	Builtins []starknetParser.Builtin
+	Builtins []sn.Builtin
 }
 
 func LoadCairoZeroProgram(content []byte) (*Program, error) {

pkg/runners/zero/zero.go

@@ -26,11 +26,11 @@ type ZeroRunner struct {
 }
 
 // Creates a new Runner of a Cairo Zero program
-func NewRunner(program *Program, proofmode bool, maxsteps uint64) (*ZeroRunner, error) {
+func NewRunner(program *Program, proofmode bool, maxsteps uint64) (ZeroRunner, error) {
 	// todo(rodro): given the program get the appropiate hints
 	hintrunner := hintrunner.NewHintRunner(make(map[uint64]hintrunner.Hinter))
 
-	return &ZeroRunner{
+	return ZeroRunner{
 		program:    program,
 		hintrunner: hintrunner,
 		proofmode:  proofmode,
@@ -76,11 +76,13 @@ func (runner *ZeroRunner) InitializeMainEntrypoint() (mem.MemoryAddress, error)
 	if runner.proofmode {
 		initialPCOffset, ok := runner.program.Labels["__start__"]
 		if !ok {
-			return mem.UnknownAddress, errors.New("start label not found. Try compiling with `--proof_mode`")
+			return mem.UnknownAddress,
+				errors.New("start label not found. Try compiling with `--proof_mode`")
 		}
 		endPcOffset, ok := runner.program.Labels["__end__"]
 		if !ok {
-			return mem.UnknownAddress, errors.New("end label not found. Try compiling with `--proof_mode`")
+			return mem.UnknownAddress,
+				errors.New("end label not found. Try compiling with `--proof_mode`")
 		}
 
 		stack := runner.initializeBuiltins(memory)
@@ -144,7 +146,9 @@ func (runner *ZeroRunner) initializeBuiltins(memory *mem.Memory) []mem.MemoryVal
 	return stack
 }
 
-func (runner *ZeroRunner) initializeVm(initialPC *mem.MemoryAddress, stack []mem.MemoryValue, memory *mem.Memory) error {
+func (runner *ZeroRunner) initializeVm(
+	initialPC *mem.MemoryAddress, stack []mem.MemoryValue, memory *mem.Memory,
+) error {
 	executionSegment := memory.Segments[vm.ExecutionSegment]
 	offset := executionSegment.Len()
 	for idx := range stack {
@@ -195,7 +199,7 @@ func (runner *ZeroRunner) RunFor(steps uint64) error {
 		if err := runner.vm.RunStep(runner.hintrunner); err != nil {
 			return fmt.Errorf(
 				"pc %s step %d: %w",
-				runner.pc().String(),
+				runner.pc(),
 				runner.steps(),
 				err,
 			)
@@ -229,16 +233,16 @@ func (runner *ZeroRunner) Output() []*fp.Element {
 	}
 
 	output := []*fp.Element{}
-	for _, segment := range runner.vm.Memory.Segments {
-		if segment.BuiltinRunner.String() == "output" {
-			for offset := uint64(0); offset < segment.Len(); offset++ {
-				value := segment.Peek(offset)
-				// todo(rodro): check if output can only contains field elements
-				valueFelt, _ := value.FieldElement()
-				output = append(output, valueFelt)
-			}
-			break
-		}
+	outputSegment, ok := runner.vm.Memory.FindSegmentWithBuiltin("output")
+	if !ok {
+		return output
+	}
+
+	for offset := uint64(0); offset < outputSegment.Len(); offset++ {
+		value := outputSegment.Peek(offset)
+		// todo(rodro): check if output can only contains field elements
+		valueFelt, _ := value.FieldElement()
+		output = append(output, valueFelt)
 	}
 	return output
 }

pkg/runners/zero/zero_test.go

@@ -1,18 +1,22 @@
 package zero
 
 import (
+	"fmt"
 	"math"
 	"testing"
 
 	"github.com/NethermindEth/cairo-vm-go/pkg/assembler"
+	sn "github.com/NethermindEth/cairo-vm-go/pkg/parsers/starknet"
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm"
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm/memory"
+	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
+	pedersenhash "github.com/consensys/gnark-crypto/ecc/stark-curve/pedersen-hash"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestSimpleProgram(t *testing.T) {
-	program := createDefaultProgram(`
+	program := createProgram(`
         [ap] = 2, ap++;
         [ap] = 3, ap++;
         [ap] = 4, ap++;
@@ -57,7 +61,7 @@ func TestSimpleProgram(t *testing.T) {
 }
 
 func TestStepLimitExceeded(t *testing.T) {
-	program := createDefaultProgram(`
+	program := createProgram(`
         [ap] = 2;
         [ap + 1] = 3;
         [ap + 2] = 5;
@@ -106,7 +110,7 @@ func TestStepLimitExceeded(t *testing.T) {
 }
 
 func TestStepLimitExceededProofMode(t *testing.T) {
-	program := createDefaultProgram(`
+	program := createProgram(`
         [ap] = 2;
         [ap + 1] = 3;
         [ap + 2] = 5;
@@ -161,6 +165,184 @@ func TestStepLimitExceededProofMode(t *testing.T) {
 	}
 }
 
+func TestBitwiseBuiltin(t *testing.T) {
+	// bitwise segment ptr is located at fp - 3 (fp - 2 and fp - 1 contain initialization vals)
+	// We first write 16 and 8 to bitwise. Then we read the bitwise result from &, ^ and |
+	runner := createRunner(`
+        [ap] = 14, ap++;
+        [ap] = 7, ap++;
+        [ap - 2] = [[fp - 3]];
+        [ap - 1] = [[fp - 3] + 1];
+
+        [ap] = [[fp - 3] + 2];
+        [ap + 1] = [[fp - 3] + 3];
+        [ap + 2] = [[fp - 3] + 4];
+
+        [ap] = 6;
+        [ap + 1] = 9;
+        [ap + 2] = 15;
+        ret;
+    `, sn.Bitwise)
+
+	err := runner.Run()
+	require.NoError(t, err)
+
+	bitwise, ok := runner.vm.Memory.FindSegmentWithBuiltin("bitwise")
+	require.True(t, ok)
+
+	requireEqualSegments(t, createSegment(14, 7, 6, 9, 15), bitwise)
+}
+
+func TestBitwiseBuiltinError(t *testing.T) {
+	// inferring first write to cell
+	runner := createRunner(`
+	    [ap] = [[fp - 3]];
+	    ret;
+	`, sn.Bitwise)
+
+	err := runner.Run()
+	require.ErrorContains(t, err, "cannot infer value")
+
+	// inferring second write to cell
+	runner = createRunner(`
+	    [ap] = [[fp - 3] + 1];
+	    ret;
+	`, sn.Bitwise)
+	err = runner.Run()
+	require.ErrorContains(t, err, "cannot infer value")
+
+	// trying to infer without writing before
+	runner = createRunner(`
+        [ap] = [[fp - 3] + 2];
+        ret;
+    `, sn.Bitwise)
+
+	err = runner.Run()
+	require.ErrorContains(t, err, "input value at offset 0 is unknown")
+}
+
+func TestOutputBuiltin(t *testing.T) {
+	// Output builtin is located at fp - 3
+	runner := createRunner(`
+        [ap] = 5;
+        [ap] = [[fp - 3]];
+        [ap + 1] = 7;
+        [ap + 1] = [[fp - 3] + 1];
+        ret;
+    `, sn.Output)
+	err := runner.Run()
+	require.NoError(t, err)
+
+	output := runner.Output()
+
+	val1 := fp.NewElement(5)
+	val2 := fp.NewElement(7)
+	require.Equal(t, []*fp.Element{&val1, &val2}, output)
+}
+
+func TestPedersenBuiltin(t *testing.T) {
+	val1 := fp.NewElement(5)
+	val2 := fp.NewElement(7)
+	val3 := pedersenhash.Pedersen(&val1, &val2)
+
+	// pedersen builtin is located at fp - 3
+	// we first write val1 and val2 and then check the infered value is val3
+	code := fmt.Sprintf(`
+        [ap] = %s;
+        [ap] = [[fp - 3]];
+
+        [ap + 1] = %s;
+        [ap + 1] = [[fp - 3] + 1];
+
+        [ap + 2] = [[fp - 3] + 2];
+        [ap + 2] = %s;
+        ret;
+    `, val1.Text(10), val2.Text(10), val3.Text(10))
+
+	runner := createRunner(code, sn.Pedersen)
+	err := runner.Run()
+	require.NoError(t, err)
+
+	pedersen, ok := runner.vm.Memory.FindSegmentWithBuiltin("pedersen")
+	require.True(t, ok)
+	requireEqualSegments(t, createSegment(&val1, &val2, &val3), pedersen)
+}
+
+func TestPedersenBuiltinError(t *testing.T) {
+	runner := createRunner(`
+        [ap] = [[fp - 3]];
+        ret;
+    `, sn.Pedersen)
+	err := runner.Run()
+	require.ErrorContains(t, err, "cannot infer value")
+
+	runner = createRunner(`
+        [ap] = [[fp - 3] + 2];
+        ret;
+    `, sn.Pedersen)
+	err = runner.Run()
+	require.ErrorContains(t, err, "input value at offset 0 is unknown")
+}
+
+func TestRangeCheckBuiltin(t *testing.T) {
+	// range check is located at fp - 3 (fp - 2 and fp - 1 contain initialization vals)
+	// we write 5 and 2**128 - 1 to range check
+	// no error should come from this
+	runner := createRunner(`
+        [ap] = 5;
+        [ap] = [[fp - 3]];
+        [ap + 1] = 0xffffffffffffffffffffffffffffffff;
+        [ap + 1] = [[fp - 3] + 1];
+        ret;
+    `, sn.RangeCheck)
+
+	err := runner.Run()
+	require.NoError(t, err)
+
+	rangeCheck, ok := runner.vm.Memory.FindSegmentWithBuiltin("range_check")
+	require.True(t, ok)
+
+	felt := &fp.Element{}
+	felt, err = felt.SetString("0xffffffffffffffffffffffffffffffff")
+	require.NoError(t, err)
+
+	requireEqualSegments(t, createSegment(5, felt), rangeCheck)
+}
+
+func TestRangeCheckBuiltinError(t *testing.T) {
+	// first test fails due to out of bound check
+	runner := createRunner(`
+        [ap] = 0x100000000000000000000000000000000;
+        [ap] = [[fp - 3]];
+        ret;
+    `, sn.RangeCheck)
+
+	err := runner.Run()
+	require.ErrorContains(t, err, "check write: 2**128 <")
+
+	// second test fails due to reading unknown value
+	runner = createRunner(`
+        [ap] = [[fp - 3]];
+        ret;
+    `, sn.RangeCheck)
+
+	err = runner.Run()
+	require.ErrorContains(t, err, "cannot infer value")
+
+}
+
+func createRunner(code string, builtins ...sn.Builtin) ZeroRunner {
+	program := createProgramWithBuiltins(code, builtins...)
+
+	runner, err := NewRunner(program, false, math.MaxUint64)
+	if err != nil {
+		panic(err)
+	}
+	return runner
+
+}
+
+// utility to create segments easier
 func createSegment(values ...any) *memory.Segment {
 	data := make([]memory.MemoryValue, len(values))
 	for i := range values {
@@ -180,14 +362,25 @@ func createSegment(values ...any) *memory.Segment {
 	return s
 }
 
+// compare two segments ignoring builtins
+func requireEqualSegments(t *testing.T, expected, result *memory.Segment) {
+	result = trimmedSegment(result)
+
+	t.Log(expected)
+	t.Log(result)
+
+	assert.Equal(t, expected.LastIndex, result.LastIndex)
+	require.Equal(t, expected.Data, result.Data)
+}
+
 // modifies a segment in place to reduce its real length to
 // its effective lenth. It returns the same segment
 func trimmedSegment(segment *memory.Segment) *memory.Segment {
 	segment.Data = segment.Data[0:segment.Len()]
 	return segment
 }
 
-func createDefaultProgram(code string) *Program {
+func createProgram(code string) *Program {
 	bytecode, err := assembler.CasmToBytecode(code)
 	if err != nil {
 		panic(err)
@@ -202,3 +395,9 @@ func createDefaultProgram(code string) *Program {
 
 	return &program
 }
+
+func createProgramWithBuiltins(code string, builtins ...sn.Builtin) *Program {
+	program := createProgram(code)
+	program.Builtins = builtins
+	return program
+}

pkg/safemath/safemath.go

@@ -46,7 +46,7 @@ func SafeOffset(x uint64, y int16) (res uint64, isOverflow bool) {
 func NextPowerOfTwo(n uint64) uint64 {
 	// it is already a power of 2
 	if (n & (n - 1)) == 0 {
-		return uint64(n)
+		return n
 	}
 
 	higherBit := 64 - bits.LeadingZeros64(n)