Add air functionality

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