Add air functionality (#647)

* Add basic air public input technology * Iterate on air private input * Iterate on air private input with bitwise * Iterate on air private input with poseidon * Iterate on air private input with pedersen * Iterate on air private input with ecop * Iterate on air private input with poseidon * Clean up empty json response * Iterate on air private input with ecdsa * Some refactoring * Add TODO * Fix unit test * Refactor air private functionality to respective builtin files * Minor updates --------- Co-authored-by: MaksymMalicki <[email protected]>
NethermindEth · Sep 12, 2024 · 636d5eb · 636d5eb
1 parent 4d93227
commit 636d5eb
Show file tree

Hide file tree

Showing 10 changed files with 539 additions and 5 deletions.
diff --git a/cmd/cli/main.go b/cmd/cli/main.go
@@ -1,9 +1,11 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"math"
 	"os"
+	"path/filepath"
 
 	"github.com/NethermindEth/cairo-vm-go/pkg/hintrunner/core"
 	"github.com/NethermindEth/cairo-vm-go/pkg/hintrunner/hinter"
@@ -24,6 +26,8 @@ func main() {
 	var traceLocation string
 	var memoryLocation string
 	var layoutName string
+	var airPublicInputLocation string
+	var airPrivateInputLocation string
 	app := &cli.App{
 		Name:                 "cairo-vm",
 		Usage:                "A cairo virtual machine",
@@ -85,6 +89,18 @@ func main() {
 						Required:    false,
 						Destination: &layoutName,
 					},
+					&cli.StringFlag{
+						Name:        "air_public_input",
+						Usage:       "location to store the air_public_input",
+						Required:    false,
+						Destination: &airPublicInputLocation,
+					},
+					&cli.StringFlag{
+						Name:        "air_private_input",
+						Usage:       "location to store the air_private_input",
+						Required:    false,
+						Destination: &airPrivateInputLocation,
+					},
 				},
 				Action: func(ctx *cli.Context) error {
 					// TODO: move this action's body to a separate function to decrease the
@@ -179,6 +195,46 @@ func main() {
 						}
 					}
 
+					if proofmode {
+						if airPublicInputLocation != "" {
+							airPublicInput, err := runner.GetAirPublicInput()
+							if err != nil {
+								return err
+							}
+							airPublicInputJson, err := json.MarshalIndent(airPublicInput, "", "  ")
+							if err != nil {
+								return err
+							}
+							err = os.WriteFile(airPublicInputLocation, airPublicInputJson, 0644)
+							if err != nil {
+								return fmt.Errorf("cannot write air_public_input: %w", err)
+							}
+						}
+
+						if airPrivateInputLocation != "" {
+							tracePath, err := filepath.Abs(traceLocation)
+							if err != nil {
+								return err
+							}
+							memoryPath, err := filepath.Abs(memoryLocation)
+							if err != nil {
+								return err
+							}
+							airPrivateInput, err := runner.GetAirPrivateInput(tracePath, memoryPath)
+							if err != nil {
+								return err
+							}
+							airPrivateInputJson, err := json.MarshalIndent(airPrivateInput, "", "  ")
+							if err != nil {
+								return err
+							}
+							err = os.WriteFile(airPrivateInputLocation, airPrivateInputJson, 0644)
+							if err != nil {
+								return fmt.Errorf("cannot write air_private_input: %w", err)
+							}
+						}
+					}
+
 					fmt.Println("Success!")
 					output := runner.Output()
 					if len(output) > 0 {

diff --git a/integration_tests/cairo_zero_file_tests/keccak_builtin.starknet_with_keccak.cairo b/integration_tests/cairo_zero_file_tests/keccak_builtin.starknet_with_keccak.cairo
@@ -0,0 +1,18 @@
+%builtins keccak
+from starkware.cairo.common.cairo_builtins import KeccakBuiltin
+from starkware.cairo.common.keccak_state import KeccakBuiltinState
+
+func main{keccak_ptr: KeccakBuiltin*}() {
+    assert keccak_ptr[0].input = KeccakBuiltinState(1, 2, 3, 4, 5, 6, 7, 8);
+    let result = keccak_ptr[0].output;
+    let keccak_ptr = keccak_ptr + KeccakBuiltin.SIZE;
+    assert result.s0 = 528644516554364142278482415480021626364691973678134577961206;
+    assert result.s1 = 768681319646568210457759892191562701823009052229295869963057;
+    assert result.s2 = 1439835513376369408063324968379272676079109225238241190228026;
+    assert result.s3 = 1150396629165612276474514703759718478742374517669870754478270;
+    assert result.s4 = 1515147102575186161827863034255579930572231617017100845406254;
+    assert result.s5 = 1412568161597072838250338588041800080889949791225997426843744;
+    assert result.s6 = 982235455376248641031519404605670648838699214888770304613539;
+    assert result.s7 = 1339947803093378278438908448344904300127577306141693325151040;
+    return ();
+}
diff --git a/pkg/runner/air_input.go b/pkg/runner/air_input.go
@@ -0,0 +1,114 @@
+package runner
+
+import (
+	"github.com/NethermindEth/cairo-vm-go/pkg/vm/builtins"
+)
+
+func (runner *ZeroRunner) GetAirPublicInput() (AirPublicInput, error) {
+	rcMin, rcMax := runner.getPermRangeCheckLimits()
+
+	// TODO: refactor to reuse earlier computed relocated trace
+	relocatedTrace := runner.vm.RelocateTrace()
+	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}
+
+	return AirPublicInput{
+		Layout:        runner.layout.Name,
+		RcMin:         rcMin,
+		RcMax:         rcMax,
+		NSteps:        len(runner.vm.Trace),
+		DynamicParams: nil,
+		// TODO: yet to be implemented fully
+		MemorySegments: memorySegments,
+		// TODO: yet to be implemented
+		PublicMemory: make([]AirPublicMemoryEntry, 0),
+	}, nil
+}
+
+type AirPublicInput struct {
+	Layout         string                           `json:"layout"`
+	RcMin          uint16                           `json:"rc_min"`
+	RcMax          uint16                           `json:"rc_max"`
+	NSteps         int                              `json:"n_steps"`
+	DynamicParams  interface{}                      `json:"dynamic_params"`
+	MemorySegments map[string]AirMemorySegmentEntry `json:"memory_segments"`
+	PublicMemory   []AirPublicMemoryEntry           `json:"public_memory"`
+}
+
+type AirMemorySegmentEntry struct {
+	BeginAddr uint64 `json:"begin_addr"`
+	StopPtr   uint64 `json:"stop_ptr"`
+}
+
+type AirPublicMemoryEntry struct {
+	Address uint16 `json:"address"`
+	Value   string `json:"value"`
+	Page    uint16 `json:"page"`
+}
+
+func (runner *ZeroRunner) GetAirPrivateInput(tracePath, memoryPath string) (AirPrivateInput, error) {
+	airPrivateInput := AirPrivateInput{
+		TracePath:  tracePath,
+		MemoryPath: memoryPath,
+	}
+
+	for _, bRunner := range runner.layout.Builtins {
+		builtinName := bRunner.Runner.String()
+		builtinSegment, ok := runner.vm.Memory.FindSegmentWithBuiltin(builtinName)
+		if ok {
+			// some checks might be missing here
+			switch builtinName {
+			case "range_check":
+				{
+					airPrivateInput.RangeCheck = bRunner.Runner.(*builtins.RangeCheck).GetAirPrivateInput(builtinSegment)
+				}
+			case "bitwise":
+				{
+					airPrivateInput.Bitwise = bRunner.Runner.(*builtins.Bitwise).GetAirPrivateInput(builtinSegment)
+				}
+			case "poseidon":
+				{
+					airPrivateInput.Poseidon = bRunner.Runner.(*builtins.Poseidon).GetAirPrivateInput(builtinSegment)
+				}
+			case "pedersen":
+				{
+					airPrivateInput.Pedersen = bRunner.Runner.(*builtins.Pedersen).GetAirPrivateInput(builtinSegment)
+				}
+			case "ec_op":
+				{
+					airPrivateInput.EcOp = bRunner.Runner.(*builtins.EcOp).GetAirPrivateInput(builtinSegment)
+				}
+			case "keccak":
+				{
+					airPrivateInput.Keccak = bRunner.Runner.(*builtins.Keccak).GetAirPrivateInput(builtinSegment)
+				}
+			case "ecdsa":
+				{
+					ecdsaAirPrivateInput, err := bRunner.Runner.(*builtins.ECDSA).GetAirPrivateInput(builtinSegment)
+					if err != nil {
+						return AirPrivateInput{}, err
+					}
+					airPrivateInput.Ecdsa = ecdsaAirPrivateInput
+				}
+			}
+		}
+	}
+
+	return airPrivateInput, nil
+}
+
+type AirPrivateInput struct {
+	TracePath  string                                 `json:"trace_path"`
+	MemoryPath string                                 `json:"memory_path"`
+	Pedersen   []builtins.AirPrivateBuiltinPedersen   `json:"pedersen"`
+	RangeCheck []builtins.AirPrivateBuiltinRangeCheck `json:"range_check"`
+	Ecdsa      []builtins.AirPrivateBuiltinECDSA      `json:"ecdsa"`
+	Bitwise    []builtins.AirPrivateBuiltinBitwise    `json:"bitwise"`
+	EcOp       []builtins.AirPrivateBuiltinEcOp       `json:"ec_op"`
+	Keccak     []builtins.AirPrivateBuiltinKeccak     `json:"keccak"`
+	Poseidon   []builtins.AirPrivateBuiltinPoseidon   `json:"poseidon"`
+}
diff --git a/pkg/vm/builtins/bitwise.go b/pkg/vm/builtins/bitwise.go
@@ -3,6 +3,8 @@ package builtins
 import (
 	"errors"
 	"fmt"
+	"math/big"
+	"sort"
 
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm/memory"
 	"github.com/consensys/gnark-crypto/ecc/stark-curve/fp"
@@ -100,3 +102,50 @@ func (b *Bitwise) String() string {
 func (b *Bitwise) GetAllocatedSize(segmentUsedSize uint64, vmCurrentStep uint64) (uint64, error) {
 	return getBuiltinAllocatedSize(segmentUsedSize, vmCurrentStep, b.ratio, inputCellsPerBitwise, instancesPerComponentBitwise, cellsPerBitwise)
 }
+
+type AirPrivateBuiltinBitwise struct {
+	Index int    `json:"index"`
+	X     string `json:"x"`
+	Y     string `json:"y"`
+}
+
+func (b *Bitwise) GetAirPrivateInput(bitwiseSegment *memory.Segment) []AirPrivateBuiltinBitwise {
+	valueMapping := make(map[int]AirPrivateBuiltinBitwise)
+	for index, value := range bitwiseSegment.Data {
+		if !value.Known() {
+			continue
+		}
+		idx, typ := index/cellsPerBitwise, index%cellsPerBitwise
+		if typ >= 2 {
+			continue
+		}
+
+		builtinValue, exists := valueMapping[idx]
+		if !exists {
+			builtinValue = AirPrivateBuiltinBitwise{Index: idx}
+		}
+
+		valueBig := big.Int{}
+		value.Felt.BigInt(&valueBig)
+		valueHex := fmt.Sprintf("0x%x", &valueBig)
+		if typ == 0 {
+			builtinValue.X = valueHex
+		} else {
+			builtinValue.Y = valueHex
+		}
+		valueMapping[idx] = builtinValue
+	}
+
+	values := make([]AirPrivateBuiltinBitwise, 0)
+
+	sortedIndexes := make([]int, 0, len(valueMapping))
+	for index := range valueMapping {
+		sortedIndexes = append(sortedIndexes, index)
+	}
+	sort.Ints(sortedIndexes)
+	for _, index := range sortedIndexes {
+		value := valueMapping[index]
+		values = append(values, value)
+	}
+	return values
+}
diff --git a/pkg/vm/builtins/ecdsa.go b/pkg/vm/builtins/ecdsa.go
@@ -2,6 +2,7 @@ package builtins
 
 import (
 	"fmt"
+	"math/big"
 
 	"github.com/NethermindEth/cairo-vm-go/pkg/utils"
 	"github.com/NethermindEth/cairo-vm-go/pkg/vm/memory"
@@ -17,7 +18,7 @@ const cellsPerECDSA = 2
 const instancesPerComponentECDSA = 1
 
 type ECDSA struct {
-	signatures map[uint64]ecdsa.Signature
+	Signatures map[uint64]ecdsa.Signature
 	ratio      uint64
 }
 
@@ -58,7 +59,7 @@ func (e *ECDSA) CheckWrite(segment *memory.Segment, offset uint64, value *memory
 	}
 
 	pubKey := &ecdsa.PublicKey{A: key}
-	sig, ok := e.signatures[pubOffset]
+	sig, ok := e.Signatures[pubOffset]
 	if !ok {
 		return fmt.Errorf("signature is missing from ECDSA builtin")
 	}
@@ -117,8 +118,8 @@ Hint that will call this function looks like this:
 	},
 */
 func (e *ECDSA) AddSignature(pubOffset uint64, r, s *fp.Element) error {
-	if e.signatures == nil {
-		e.signatures = make(map[uint64]ecdsa.Signature)
+	if e.Signatures == nil {
+		e.Signatures = make(map[uint64]ecdsa.Signature)
 	}
 	bytes := make([]byte, 0, 64)
 	rBytes := r.Bytes()
@@ -132,7 +133,7 @@ func (e *ECDSA) AddSignature(pubOffset uint64, r, s *fp.Element) error {
 		return err
 	}
 
-	e.signatures[pubOffset] = sig
+	e.Signatures[pubOffset] = sig
 	return nil
 }
 
@@ -162,3 +163,49 @@ func recoverY(x *fp.Element) (fp.Element, fp.Element, error) {
 	negY.Neg(y)
 	return *y, negY, nil
 }
+
+type AirPrivateBuiltinECDSASignatureInput struct {
+	R string `json:"r"`
+	W string `json:"w"`
+}
+
+type AirPrivateBuiltinECDSA struct {
+	Index          int                                  `json:"index"`
+	PubKey         string                               `json:"pubkey"`
+	Msg            string                               `json:"msg"`
+	SignatureInput AirPrivateBuiltinECDSASignatureInput `json:"signature_input"`
+}
+
+func (e *ECDSA) GetAirPrivateInput(ecdsaSegment *memory.Segment) ([]AirPrivateBuiltinECDSA, error) {
+	values := make([]AirPrivateBuiltinECDSA, 0)
+	for addrOffset, signature := range e.Signatures {
+		idx := addrOffset / cellsPerECDSA
+		pubKey, err := ecdsaSegment.Read(addrOffset)
+		if err != nil {
+			return values, err
+		}
+		msg, err := ecdsaSegment.Read(addrOffset + 1)
+		if err != nil {
+			return values, err
+		}
+
+		pubKeyBig := big.Int{}
+		msgBig := big.Int{}
+		pubKey.Felt.BigInt(&pubKeyBig)
+		msg.Felt.BigInt(&msgBig)
+		pubKeyHex := fmt.Sprintf("0x%x", &pubKeyBig)
+		msgHex := fmt.Sprintf("0x%x", &msgBig)
+
+		rBig := new(big.Int).SetBytes(signature.R[:])
+		sBig := new(big.Int).SetBytes(signature.S[:])
+		frModulusBig, _ := new(big.Int).SetString("3618502788666131213697322783095070105526743751716087489154079457884512865583", 10)
+		wBig := new(big.Int).ModInverse(sBig, frModulusBig)
+		signatureInput := AirPrivateBuiltinECDSASignatureInput{
+			R: fmt.Sprintf("0x%x", rBig),
+			W: fmt.Sprintf("0x%x", wBig),
+		}
+
+		values = append(values, AirPrivateBuiltinECDSA{Index: int(idx), PubKey: pubKeyHex, Msg: msgHex, SignatureInput: signatureInput})
+	}
+	return values, nil
+}