feat: parse abi (#2)

* feat: parse abi * chore: update README.md * chore: update readme --------- Co-authored-by: Yi Sun <[email protected]>
axiom-crypto · Feb 9, 2024 · fcad24b · fcad24b
1 parent 2876edd
commit fcad24b
Show file tree

Hide file tree

Showing 13 changed files with 377 additions and 165 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -24,12 +24,12 @@ jobs:
         run: |
           npm install
           mv build/axiom-std-cli-build.js build/axiom-std-cli-temp-build.js
-          cp src/test/AxiomCli.sol src/test/AxiomCli.sol.temp
+          cp src/AxiomCli.sol src/AxiomCli.sol.temp
           npm run build
           diff build/axiom-std-cli-build.js build/axiom-std-cli-temp-build.js
-          diff src/test/AxiomCli.sol src/test/AxiomCli.sol.temp
+          diff src/AxiomCli.sol src/AxiomCli.sol.temp
           rm build/axiom-std-cli-temp-build.js
-          rm src/test/AxiomCli.sol.temp
+          rm src/AxiomCli.sol.temp
 
       - name: Install Foundry
         uses: foundry-rs/foundry-toolchain@v1
@@ -61,12 +61,12 @@ jobs:
         run: |
           npm install
           mv build/axiom-std-cli-build.js build/axiom-std-cli-temp-build.js
-          cp src/test/AxiomCli.sol src/test/AxiomCli.sol.temp
+          cp src/AxiomCli.sol src/AxiomCli.sol.temp
           npm run build
           diff build/axiom-std-cli-build.js build/axiom-std-cli-temp-build.js
-          diff src/test/AxiomCli.sol src/test/AxiomCli.sol.temp
+          diff src/AxiomCli.sol src/AxiomCli.sol.temp
           rm build/axiom-std-cli-temp-build.js
-          rm src/test/AxiomCli.sol.temp
+          rm src/AxiomCli.sol.temp
 
       - name: Install Foundry
         uses: foundry-rs/foundry-toolchain@v1

diff --git a/README.md b/README.md
@@ -1,16 +1,16 @@
-## Axiom V2 Periphery
+## axiom-std
 
-This repo contains client smart contracts, interfaces, and testing utilities for applications integrating Axiom V2. To learn more about how to integrate Axiom into your application, see the [developer docs](https://docs.axiom.xyz). For the complete Axiom V2 smart contract code, see the smart contract repo [here](https://github.com/axiom-crypto/axiom-v2-contracts).
+This repo contains Foundry testing utilities for applications integrating Axiom V2. To learn more about how to integrate Axiom into your application, see the [developer docs](https://docs.axiom.xyz). For the complete Axiom V2 smart contract code, see the smart contract repo [here](https://github.com/axiom-crypto/axiom-v2-contracts).
 
 ## Installation
 
 To use smart contracts or test utilities from this repo in your **external Foundry project**, run:
 
 ```bash
-forge install axiom-crypto/axiom-v2-periphery
+forge install axiom-crypto/axiom-std
 ```
 
-Add `@axiom-crypto/v2-periphery/=lib/axiom-v2-periphery/src` in `remappings.txt`.
+Add `@axiom-crypto/axiom-std/=lib/axiom-std/src` in `remappings.txt`.
 
 ## Usage
 
@@ -65,102 +65,57 @@ contract AverageBalance is AxiomV2Client {
 
 #### Testing with `AxiomTest` Foundry tests
 
-To test your code, you can use `AxiomTest.sol` in place of `forge-std/Test.sol`. This extension to the standard Foundry test library provides Axiom-specific cheatcodes accessible to your Foundry tests. Using these cheatcodes requires the Axiom Client SDK, which is provided via the npm package `@axiom-crypto/client`; you can install this in your Foundry project using
+To test your code, you can use `AxiomTest.sol` in place of `forge-std/Test.sol`. This extension to the standard Foundry test library provides Axiom-specific cheatcodes accessible to your Foundry tests. Using these cheatcodes requires Node to be installed.
 
-```bash
-npm install @axiom-crypto/client
-yarn add @axiom-crypto/client
-pnpm add @axiom-crypto/client
-```
-
-Once you have written an Axiom circuit, you can test it against your client smart contract using the `AxiomVm` cheatcodes `sendQueryArgs`, `fulfillCallbackArgs`, `prankCallback`, and `prankOffchainCallback`:
+Once you have written an Axiom circuit, you can test it against your client smart contract by specifying your circuit input struct `AxiomInput`, constructing a `Query`, and then using the `Query.send()` and `Query.prankFulfill()` cheatcodes:
 
 ```solidity
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.19;
 
-import { AxiomTest, AxiomVm } from "../src/test/AxiomTest.sol";
-import { IAxiomV2Query } from "../src/interfaces/query/IAxiomV2Query.sol";
+import "@axiom-crypto/axiom-std/AxiomTest.sol";
 
 import { AverageBalance } from "./example/AverageBalance.sol";
 
 contract AverageBalanceTest is AxiomTest {
+    using Axiom for Query;
+
+    struct AxiomInput {
+        uint64 blockNumber;
+        address address;
+    }
+
     AverageBalance public averageBalance;
+    AxiomInput public defaultInput;
+    bytes32 public querySchema;
 
     function setUp() public {
-        _createSelectForkAndSetupAxiom("sepolia", 5_057_320);
+        _createSelectForkAndSetupAxiom("sepolia", 5_103_100);
 
-        inputPath = "test/circuit/input.json";
-        querySchema = axiomVm.compile("test/circuit/average.circuit.ts", inputPath);
+        defaultInput =
+            AxiomInput({ blockNumber: 4_205_938, _address: address(0x8018fe32fCFd3d166E8b4c4E37105318A84BA11b) });
+        querySchema = axiomVm.readCircuit("test/circuit/average.circuit.ts", abi.encode(defaultInput));
         averageBalance = new AverageBalance(axiomV2QueryAddress, uint64(block.chainid), querySchema);
     }
 
-    function test_axiomSendQuery() public {
-        axiomVm.getArgsAndSendQuery(
-            inputPath,
-            address(averageBalance),
-            callbackExtraData,
-            feeData,
-            msg.sender
-        );
-    }
-
-    function test_axiomSendQueryWithArgs() public {
-        AxiomVm.AxiomSendQueryArgs memory args = axiomVm.sendQueryArgs(
-            inputPath, address(averageBalance), callbackExtraData, feeData
-        );
-        axiomV2Query.sendQuery{ value: args.value }(
-            args.sourceChainId,
-            args.dataQueryHash,
-            args.computeQuery,
-            args.callback,
-            args.feeData,
-            args.userSalt,
-            args.refundee,
-            args.dataQuery
-        );
-    }
+    function test_simple_example() public {
+        // create a query into Axiom with default parameters
+        Query memory q = query(querySchema, abi.encode(defaultInput), address(averageBalance));
 
-    function test_axiomCallback() public {
-        axiomVm.prankCallback(
-            inputPath,
-            address(averageBalance),
-            callbackExtraData,
-            feeData,
-            msg.sender
-        );
-    }
+        // send the query to Axiom
+        q.send();
 
-    function test_AxiomCallbackWithArgs() public {
-        AxiomVm.AxiomFulfillCallbackArgs memory args = axiomVm.fulfillCallbackArgs(
-            inputPath,
-            address(averageBalance),
-            callbackExtraData,
-            feeData,
-            msg.sender
-        );
-        axiomVm.prankCallback(args);
-    }
+        // prank fulfillment of the query, returning the Axiom results 
+        bytes32[] memory results = q.prankFulfill();
 
-    function test_axiomOffchainCallback() public {
-        axiomVm.prankOffchainCallback(
-            inputPath,
-            address(averageBalance),
-            callbackExtraData,
-            feeData,
-            msg.sender
-        );
-    }
+        // parse Axiom results and verify length is as expected
+        assertEq(results.length, 3);
+        uint256 blockNumber = uint256(results[0]);
+        address addr = address(uint160(uint256(results[1])));
+        uint256 avg = uint256(results[2]);
 
-    function test_AxiomOffchainCallbackWithArgs() public {
-        AxiomVm.AxiomFulfillCallbackArgs memory args = axiomVm.fulfillCallbackArgs(
-            inputPath,
-            address(averageBalance),
-            callbackExtraData,
-            feeData,
-            msg.sender
-        );
-        axiomVm.prankOffchainCallback(args);
+        // verify the average balance recorded in AverageBalance is as expected
+        assertEq(avg, averageBalance.provenAverageBalances(blockNumber, addr));
     }
 }
 
@@ -172,5 +127,5 @@ This repo contains both Foundry and Javascript packages. To install, run:
 
 ```bash
 forge install
-pnpm install     # or `npm install` or `yarn install`
+npm install     # or `pnpm install` or `yarn install`
 ```
diff --git a/build/axiom-std-cli-build.js b/build/axiom-std-cli-build.js
@@ -9,7 +9,11 @@ var require_utils = __commonJS({
   "dist/utils.js"(exports2) {
     "use strict";
     Object.defineProperty(exports2, "__esModule", { value: true });
-    exports2.getSolidityType = void 0;
+    exports2.getInputs = exports2.getAbi = exports2.findStructDefinition = exports2.findFilesWithAxiomInput = exports2.getSolidityType = void 0;
+    var tslib_1 = require("tslib");
+    var fs_1 = tslib_1.__importDefault(require("fs"));
+    var path_1 = tslib_1.__importDefault(require("path"));
+    var viem_1 = require("viem");
     var getSolidityType = (type) => {
       switch (type) {
         case "CircuitValue":
@@ -25,6 +29,80 @@ var require_utils = __commonJS({
       }
     };
     exports2.getSolidityType = getSolidityType;
+    var findFilesWithAxiomInput = (directory) => {
+      let file = null;
+      function traverseDirectory(dir) {
+        const entries = fs_1.default.readdirSync(dir);
+        for (const entry of entries) {
+          const entryPath = path_1.default.join(dir, entry);
+          const stat = fs_1.default.statSync(entryPath);
+          if (stat.isDirectory()) {
+            traverseDirectory(entryPath);
+          } else if (stat.isFile() && entry.endsWith(".json")) {
+            const fileContent = fs_1.default.readFileSync(entryPath, "utf8");
+            if (fileContent.includes('.AxiomInput"')) {
+              file = entryPath;
+              return;
+            }
+          }
+        }
+      }
+      traverseDirectory(directory);
+      return file;
+    };
+    exports2.findFilesWithAxiomInput = findFilesWithAxiomInput;
+    var findStructDefinition = (jsonFile) => {
+      const jsonData = require(jsonFile);
+      const fileName = path_1.default.basename(jsonFile, path_1.default.extname(jsonFile));
+      function traverseObject(obj) {
+        if (obj.nodeType === "StructDefinition" && obj.canonicalName === `${fileName}.AxiomInput`) {
+          return obj;
+        }
+        for (const key in obj) {
+          if (typeof obj[key] === "object") {
+            const result = traverseObject(obj[key]);
+            if (result !== null) {
+              return result;
+            }
+          }
+        }
+        return null;
+      }
+      return traverseObject(jsonData);
+    };
+    exports2.findStructDefinition = findStructDefinition;
+    var getAbi = () => {
+      const jsonFile = (0, exports2.findFilesWithAxiomInput)(process.cwd());
+      if (jsonFile === null) {
+        throw new Error("Could not find json file with AxiomInput");
+      }
+      const structDefinition = (0, exports2.findStructDefinition)(jsonFile);
+      if (structDefinition === null) {
+        throw new Error(`Could not find struct definition in file ${jsonFile}`);
+      }
+      const abi = [];
+      for (const member of structDefinition.members) {
+        const type = member.typeDescriptions.typeString;
+        if (type === void 0) {
+          throw new Error(`Could not find type for member ${member.name}`);
+        }
+        abi.push({ name: member.name, type });
+      }
+      return abi;
+    };
+    exports2.getAbi = getAbi;
+    var getInputs = (inputs2, inputSchema) => {
+      const inputSchemaJson = JSON.parse(inputSchema);
+      const keys = Object.keys(inputSchemaJson);
+      const abi = (0, exports2.getAbi)();
+      const rawInputs = (0, viem_1.decodeAbiParameters)(abi, inputs2);
+      const circuitInputs2 = {};
+      for (let i = 0; i < keys.length; i++) {
+        circuitInputs2[keys[i]] = rawInputs[i].toString();
+      }
+      return circuitInputs2;
+    };
+    exports2.getInputs = getInputs;
   }
 });
 
@@ -36,7 +114,6 @@ var require_compile = __commonJS({
     exports2.compile = void 0;
     var js_12 = require("@axiom-crypto/circuit/js");
     var utils_12 = require("@axiom-crypto/circuit/cliHandler/utils");
-    var viem_12 = require("viem");
     var utils_22 = require_utils();
     var compile = async (circuitPath, inputs2, providerUri2) => {
       let circuitFunction = "circuit";
@@ -49,17 +126,7 @@ var require_compile = __commonJS({
         shouldTime: false,
         inputSchema: f.inputSchema
       });
-      let abi2 = [];
-      let inputSchemaJson2 = JSON.parse(f.inputSchema);
-      let keys2 = Object.keys(inputSchemaJson2);
-      for (let i = 0; i < keys2.length; i++) {
-        abi2.push({ "name": keys2[i], "type": (0, utils_22.getSolidityType)(inputSchemaJson2[keys2[i]]) });
-      }
-      const rawInputs2 = (0, viem_12.decodeAbiParameters)(abi2, inputs2);
-      const circuitInputs2 = {};
-      for (let i = 0; i < keys2.length; i++) {
-        circuitInputs2[keys2[i]] = rawInputs2[i].toString();
-      }
+      const circuitInputs2 = (0, utils_22.getInputs)(inputs2, f.inputSchema);
       try {
         const res = await circuit2.mockCompile(circuitInputs2);
         const circuitFn = `const ${f.importName} = AXIOM_CLIENT_IMPORT
@@ -88,7 +155,6 @@ var require_prove = __commonJS({
     var js_1 = require("@axiom-crypto/circuit/js");
     var utils_1 = require("@axiom-crypto/circuit/cliHandler/utils");
     var utils_2 = require_utils();
-    var viem_1 = require("viem");
     var core_1 = require("@axiom-crypto/core");
     var client_1 = require("@axiom-crypto/client");
     var utils_3 = require("@axiom-crypto/client/axiom/utils");
@@ -102,23 +168,12 @@ var require_prove = __commonJS({
       const circuit = new js_1.AxiomBaseCircuit({
         f: eval(raw),
         mock: true,
-        // chainId: options.sourceChainId,
         provider,
         shouldTime: false,
         inputSchema: compiled.inputSchema
       });
-      let abi = [];
       let decodedInputSchema = Buffer.from(compiled.inputSchema, "base64");
-      let inputSchemaJson = JSON.parse(decoder.decode(decodedInputSchema));
-      let keys = Object.keys(inputSchemaJson);
-      for (let i = 0; i < keys.length; i++) {
-        abi.push({ "name": keys[i], "type": (0, utils_2.getSolidityType)(inputSchemaJson[keys[i]]) });
-      }
-      const rawInputs = (0, viem_1.decodeAbiParameters)(abi, inputs);
-      const circuitInputs = {};
-      for (let i = 0; i < keys.length; i++) {
-        circuitInputs[keys[i]] = rawInputs[i].toString();
-      }
+      const circuitInputs = (0, utils_2.getInputs)(inputs, decoder.decode(decodedInputSchema));
       const axiom = new core_1.AxiomSdkCore({
         providerUri: provider,
         chainId: sourceChainId,

diff --git a/cli/compile.ts b/cli/compile.ts
@@ -1,7 +1,7 @@
 import { AxiomBaseCircuit } from "@axiom-crypto/circuit/js";
 import { getFunctionFromTs, getProvider, readInputs, saveJsonToFile } from "@axiom-crypto/circuit/cliHandler/utils";
 import { decodeAbiParameters, encodeAbiParameters } from 'viem';
-import { getSolidityType } from "./utils";
+import { getAbi, getInputs, getSolidityType } from "./utils";
 
 export const compile = async (
     circuitPath: string,
@@ -19,20 +19,7 @@ export const compile = async (
         inputSchema: f.inputSchema,
     })
 
-    let abi: { name: string; type: string; }[] = [];
-    let inputSchemaJson = JSON.parse(f.inputSchema);
-    let keys = Object.keys(inputSchemaJson);
-    for (let i = 0; i < keys.length; i++) {
-        abi.push({ "name": keys[i], "type": getSolidityType(inputSchemaJson[keys[i]]) });
-    }
-    // console.log(abi);
-    // console.log(encodeAbiParameters(abi, [4205938, "0x8018fe32fCFd3d166E8b4c4E37105318A84BA11b"]))
-    const rawInputs: any[] = decodeAbiParameters(abi, inputs as `0x${string}`);
-    const circuitInputs: any = {};
-    for (let i = 0; i < keys.length; i++) {
-        circuitInputs[keys[i]] = rawInputs[i].toString();
-    }
-    // console.log(circuitInputs);
+    const circuitInputs = getInputs(inputs, f.inputSchema);
 
     try {
         const res = await circuit.mockCompile(circuitInputs);
@@ -43,9 +30,6 @@ export const compile = async (
             ...res,
             circuit: Buffer.from(circuitBuild).toString('base64'),
         }
-
-        // const finalBuild = encoder.encode(JSON.stringify(build));
-        // console.log(Buffer.from(finalBuild).toString('base64'));
         console.log(JSON.stringify(build));
     }
     catch (e) {