feat: cache Rust path

axiom-crypto · May 21, 2024 · 3389687 · 3389687
1 parent e47d118
commit 3389687
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 9 deletions.
diff --git a/src/AxiomTest.sol b/src/AxiomTest.sol
@@ -66,9 +66,7 @@ abstract contract AxiomTest is Test {
 
     /// @dev Create a forked test environment from the latest block and set up Axiom contracts
     /// @param urlOrAlias The URL or alias of the fork to create
-    function _createSelectForkAndSetupAxiom(
-        string memory urlOrAlias
-    ) internal {
+    function _createSelectForkAndSetupAxiom(string memory urlOrAlias) internal {
         vm.createSelectFork(urlOrAlias);
         _setupAxiomFromFork(block.number);
 
@@ -87,9 +85,7 @@ abstract contract AxiomTest is Test {
 
     /// @dev Set up Axiom contracts
     /// @param forkBlock The block number that the fork was created from
-    function _setupAxiomFromFork(
-        uint256 forkBlock
-    ) private {
+    function _setupAxiomFromFork(uint256 forkBlock) private {
         uint64 chainId = uint64(block.chainid);
 
         if (chainId == MAINNET_CHAIN_ID || chainId == BASE_CHAIN_ID) {

diff --git a/src/AxiomVm.sol b/src/AxiomVm.sol
@@ -122,17 +122,26 @@ library Axiom {
 /// @dev A contract that provides cheatcodes for testing the AxiomV2Query contract
 contract AxiomVm is Test {
     /// @dev Path to the Axiom CLI
-    string CLI_PATH;
+    string public CLI_PATH;
 
     /// @dev Command to run node scripts
-    string NODE_PATH;
+    string public NODE_PATH;
 
     /// @dev The URL or alias of the JSON RPC provider
-    string urlOrAlias;
+    string public urlOrAlias;
 
     address public axiomV2QueryAddress;
     mapping(bytes32 => string) compiledStrings;
 
+    /// @dev Boolean indicating if the circuit is a Rust circuit
+    bool public isRustCircuit;
+
+    /// @dev Default dummy query schema for Rust circuits
+    bytes32 constant DEFAULT_RUST_QUERY_SCHEMA = 0xdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef;
+
+    /// @dev Circuit path for Rust circuits
+    string public rustCircuitPath;
+
     constructor(address _axiomV2QueryAddress, string memory _urlOrAlias) {
         axiomV2QueryAddress = _axiomV2QueryAddress;
         urlOrAlias = _urlOrAlias;
@@ -253,6 +262,17 @@ contract AxiomVm is Test {
         compiledStrings[querySchema] = build;
     }
 
+    /**
+     * @dev Compiles a Rust circuit using the Axiom CLI via FFI
+     * @param _rustCircuitPath path to the Rust circuit file
+     * @return querySchema
+     */
+    function readRustCircuit(string memory _rustCircuitPath) public returns (bytes32 querySchema) {
+        isRustCircuit = true;
+        rustCircuitPath = _rustCircuitPath;
+        querySchema = DEFAULT_RUST_QUERY_SCHEMA;
+    }
+
     /**
      * @dev Generates args for the sendQuery function
      * @param querySchema the query schema
@@ -488,6 +508,20 @@ contract AxiomVm is Test {
         address callbackTarget,
         bytes memory callbackExtraData,
         IAxiomV2Query.AxiomV2FeeData memory feeData
+    ) internal returns (string memory output) {
+        if (!isRustCircuit) {
+            output = _runJsCircuit(querySchema, input, callbackTarget, callbackExtraData, feeData);
+        } else {
+            output = _runRustCircuit(querySchema, input, callbackTarget, callbackExtraData, feeData);
+        }
+    }
+
+    function _runJsCircuit(
+        bytes32 querySchema,
+        bytes memory input,
+        address callbackTarget,
+        bytes memory callbackExtraData,
+        IAxiomV2Query.AxiomV2FeeData memory feeData
     ) internal returns (string memory output) {
         require(bytes(compiledStrings[querySchema]).length > 0, "Circuit has not been compiled. Run `compile` first.");
         string[] memory cli = new string[](13);
@@ -512,6 +546,33 @@ contract AxiomVm is Test {
         output = build;
     }
 
+    function _runRustCircuit(
+        bytes32 querySchema,
+        bytes memory input,
+        address callbackTarget,
+        bytes memory callbackExtraData,
+        IAxiomV2Query.AxiomV2FeeData memory feeData
+    ) internal returns (string memory output) {
+        // TODO: finish this call
+        string[] memory cli = new string[](11);
+        cli[0] = "cargo";
+        cli[1] = "run";
+        cli[2] = "--circuit";
+        cli[3] = "main";
+        cli[4] = "--";
+        cli[5] = "witness-gen";
+        cli[6] = "--input";
+        cli[7] = "data/input.json";
+        cli[8] = "-k";
+        cli[9] = "12";
+        cli[10] = "-p";
+
+        bytes memory axiomOutput = vm.ffi(cli);
+        (string memory logs, string memory errors, string memory build) =
+            abi.decode(axiomOutput, (string, string, string));
+        output = build;
+    }
+
     /**
      * @dev Parses AxiomQuery from the CLI calldata bytes output
      * @param _queryString the string output from the CLI