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]>

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 {

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 ();
+}

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"`
+}

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
+}

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
+}