feat: allow multiple axiom input structs (#4)

axiom-crypto · Feb 14, 2024 · 4f9ad0e · 4f9ad0e
1 parent 089c09a
commit 4f9ad0e
Show file tree

Hide file tree

Showing 8 changed files with 192 additions and 92 deletions.
diff --git a/build/axiom-std-cli-build.js b/build/axiom-std-cli-build.js
@@ -47,28 +47,13 @@ var require_utils = __commonJS({
   "dist/utils.js"(exports2) {
     "use strict";
     Object.defineProperty(exports2, "__esModule", { value: true });
-    exports2.getInputs = exports2.getAbi = exports2.findStructDefinition = exports2.findFilesWithAxiomInput = exports2.getSolidityType = void 0;
+    exports2.getInputs = exports2.getAbis = exports2.findStructDefinition = exports2.findFilesWithAxiomInput = 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":
-          return "uint256";
-        case "CircuitValue256":
-          return "uint256";
-        case "CircuitValue[]":
-          return "uint256[]";
-        case "CircuitValue256[]":
-          return "uint256[]";
-        default:
-          throw new Error(`Unknown type ${type}`);
-      }
-    };
-    exports2.getSolidityType = getSolidityType;
     var findFilesWithAxiomInput = (directory) => {
-      let file = null;
+      let files = [];
       function traverseDirectory(dir) {
         const entries = fs_1.default.readdirSync(dir);
         for (const entry of entries) {
@@ -79,14 +64,13 @@ var require_utils = __commonJS({
           } else if (stat.isFile() && entry.endsWith(".json")) {
             const fileContent = fs_1.default.readFileSync(entryPath, "utf8");
             if (fileContent.includes('.AxiomInput"')) {
-              file = entryPath;
-              return;
+              files.push(entryPath);
             }
           }
         }
       }
       traverseDirectory(directory);
-      return file;
+      return files;
     };
     exports2.findFilesWithAxiomInput = findFilesWithAxiomInput;
     var findStructDefinition = (jsonFile) => {
@@ -109,42 +93,72 @@ var require_utils = __commonJS({
       return traverseObject(jsonData);
     };
     exports2.findStructDefinition = findStructDefinition;
-    var getAbi = () => {
-      const jsonFile = (0, exports2.findFilesWithAxiomInput)(process.cwd());
-      if (jsonFile === null) {
+    var getAbis = () => {
+      const jsonFiles = (0, exports2.findFilesWithAxiomInput)(process.cwd());
+      if (jsonFiles.length === 0) {
         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 structDefinitions = jsonFiles.map(exports2.findStructDefinition).filter((x) => x !== null);
+      if (structDefinitions.length === 0) {
+        throw new Error(`Could not find struct definition`);
       }
-      const abiComponents = [];
-      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}`);
+      const getAbiFromStructDefinition = (structDefinition) => {
+        const abiComponents = [];
+        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}`);
+          }
+          abiComponents.push({ name: member.name, type });
         }
-        abiComponents.push({ name: member.name, type });
+        const abi = [{
+          "name": "circuit",
+          "type": "tuple",
+          "components": abiComponents
+        }];
+        return abi;
+      };
+      const abis = structDefinitions.map(getAbiFromStructDefinition);
+      return abis;
+    };
+    exports2.getAbis = getAbis;
+    var validateAbi = (abi, inputSchema) => {
+      const inputSchemaJson = JSON.parse(inputSchema);
+      const keys = Object.keys(inputSchemaJson);
+      const values = Object.values(inputSchemaJson);
+      const abiComponents = abi[0].components;
+      if (keys.length !== abiComponents.length)
+        return false;
+      for (let i = 0; i < keys.length; i++) {
+        if (values[i].endsWith("[]") && !abiComponents[i].type.endsWith("[]"))
+          return false;
+        if (!values[i].endsWith("[]") && abiComponents[i].type.endsWith("[]"))
+          return false;
+        if (values[i].startsWith("CircuitValue256") && !(abiComponents[i].type.startsWith("uint256") || abiComponents[i].type.startsWith("bytes32")))
+          return false;
+        if ((abiComponents[i].type.startsWith("uint256") || abiComponents[i].type.startsWith("bytes32")) && !values[i].startsWith("CircuitValue256"))
+          return false;
+        if (!(abiComponents[i].type.startsWith("uint") || abiComponents[i].type.startsWith("address") || abiComponents[i].type.startsWith("bytes") || abiComponents[i].type.startsWith("bool")))
+          return false;
       }
-      const abi = [{
-        "name": "circuit",
-        "type": "tuple",
-        "components": abiComponents
-      }];
-      return abi;
+      return true;
     };
-    exports2.getAbi = getAbi;
     var getInputs = (inputs2, inputSchema) => {
       const inputSchemaJson = JSON.parse(inputSchema);
       const keys = Object.keys(inputSchemaJson);
-      const abi = (0, exports2.getAbi)();
+      const abis = (0, exports2.getAbis)().filter((x) => validateAbi(x, inputSchema));
+      if (abis.length === 0) {
+        throw new Error("Could not find valid ABI");
+      }
+      const abi = abis[0];
       const rawInputs = (0, viem_1.decodeAbiParameters)(abi, inputs2)[0];
+      const abiComponents = abi[0].components;
       const circuitInputs2 = {};
       for (let i = 0; i < keys.length; i++) {
         if (Array.isArray(rawInputs[keys[i]])) {
-          circuitInputs2[keys[i]] = rawInputs[keys[i]].map((x) => x.toString());
+          circuitInputs2[keys[i]] = rawInputs[abiComponents[i].name].map((x) => x.toString());
         } else {
-          circuitInputs2[keys[i]] = rawInputs[keys[i]].toString();
+          circuitInputs2[keys[i]] = rawInputs[abiComponents[i].name].toString();
         }
       }
       return circuitInputs2;

diff --git a/cli/compile.ts b/cli/compile.ts
@@ -1,7 +1,5 @@
 import { AxiomBaseCircuit } from "@axiom-crypto/circuit/js";
-import { getFunctionFromTs, getProvider, readInputs, saveJsonToFile } from "@axiom-crypto/circuit/cliHandler/utils";
-import { decodeAbiParameters, encodeAbiParameters } from 'viem';
-import { getAbi, getInputs, getSolidityType } from "./utils";
+import { getFunctionFromTs, getProvider } from "@axiom-crypto/circuit/cliHandler/utils";
 
 export const compile = async (
     circuitPath: string,

diff --git a/cli/prove.ts b/cli/prove.ts
@@ -1,7 +1,7 @@
 import path from 'path';
 import { AxiomBaseCircuit } from "@axiom-crypto/circuit/js";
 import { getFunctionFromTs, getProvider, readInputs, saveJsonToFile } from "@axiom-crypto/circuit/cliHandler/utils";
-import { getAbi, getInputs, getSolidityType } from './utils';
+import { getInputs } from './utils';
 import { decodeAbiParameters } from 'viem';
 import { AxiomSdkCore } from "@axiom-crypto/core";
 import { buildSendQuery } from "@axiom-crypto/client";

diff --git a/cli/utils.ts b/cli/utils.ts
@@ -2,23 +2,8 @@ import fs from 'fs';
 import path from 'path';
 import { decodeAbiParameters } from 'viem';
 
-export const getSolidityType = (type: string) => {
-    switch (type) {
-        case "CircuitValue":
-            return "uint256";
-        case "CircuitValue256":
-            return "uint256";
-        case "CircuitValue[]":
-            return "uint256[]";
-        case "CircuitValue256[]":
-            return "uint256[]";
-        default:
-            throw new Error(`Unknown type ${type}`);
-    }
-}
-
-export const findFilesWithAxiomInput = (directory: string): string | null => {
-    let file: string | null = null;
+export const findFilesWithAxiomInput = (directory: string): string[] => {
+    let files: string[] = [];
 
     function traverseDirectory(dir: string): void {
         const entries = fs.readdirSync(dir);
@@ -32,16 +17,15 @@ export const findFilesWithAxiomInput = (directory: string): string | null => {
             } else if (stat.isFile() && entry.endsWith('.json')) {
                 const fileContent = fs.readFileSync(entryPath, 'utf8');
                 if (fileContent.includes('.AxiomInput"')) {
-                    file = entryPath;
-                    return;
+                    files.push(entryPath);
                 }
             }
         }
     }
 
     traverseDirectory(directory);
 
-    return file;
+    return files;
 }
 
 export const findStructDefinition = (jsonFile: string): any | null => {
@@ -68,52 +52,84 @@ export const findStructDefinition = (jsonFile: string): any | null => {
     return traverseObject(jsonData);
 }
 
-export const getAbi = (): any => { 
-    const jsonFile = findFilesWithAxiomInput(process.cwd());
+export const getAbis = (): any => {
+    const jsonFiles = findFilesWithAxiomInput(process.cwd());
 
-    if (jsonFile === null) {
+    if (jsonFiles.length === 0) {
         throw new Error("Could not find json file with AxiomInput");
     }
 
-    const structDefinition = findStructDefinition(jsonFile);
+    // const structDefinition = findStructDefinition(jsonFile);
 
-    if (structDefinition === null) {
-        throw new Error(`Could not find struct definition in file ${jsonFile}`);
+    const structDefinitions = jsonFiles.map(findStructDefinition).filter((x) => x !== null);
+
+    if (structDefinitions.length === 0) {
+        throw new Error(`Could not find struct definition`);
     }
 
-    const abiComponents: { name: string; type: string; }[] = [];
+    const getAbiFromStructDefinition = (structDefinition: any) => {
+        const abiComponents: { name: string; type: string; }[] = [];
 
-    for (const member of structDefinition.members) {
-        const type = member.typeDescriptions.typeString;
-        if (type === undefined) {
-            throw new Error(`Could not find type for member ${member.name}`);
+        for (const member of structDefinition.members) {
+            const type = member.typeDescriptions.typeString;
+            if (type === undefined) {
+                throw new Error(`Could not find type for member ${member.name}`);
+            }
+            abiComponents.push({ name: member.name, type });
         }
-        abiComponents.push({ name: member.name, type });
+
+        const abi = [{
+            "name": "circuit",
+            "type": "tuple",
+            "components": abiComponents,
+        }];
+        return abi;
     }
 
-    const abi = [{
-        "name": "circuit",
-        "type": "tuple",
-        "components": abiComponents,
-    }];
-    return abi;
+    const abis = structDefinitions.map(getAbiFromStructDefinition);
+
+    return abis;
 
 }
 
+const validateAbi = (abi: any, inputSchema: string): boolean => {
+    const inputSchemaJson = JSON.parse(inputSchema);
+    const keys = Object.keys(inputSchemaJson);
+    const values: string[] = Object.values(inputSchemaJson);
+    const abiComponents = abi[0].components;
+
+    if (keys.length !== abiComponents.length) return false;
+    for (let i = 0; i < keys.length; i++) {
+        if (values[i].endsWith('[]') && !abiComponents[i].type.endsWith('[]')) return false;
+        if (!values[i].endsWith('[]') && abiComponents[i].type.endsWith('[]')) return false;
+        if (values[i].startsWith('CircuitValue256') && !(abiComponents[i].type.startsWith('uint256') || abiComponents[i].type.startsWith('bytes32'))) return false;
+        if ((abiComponents[i].type.startsWith('uint256') || abiComponents[i].type.startsWith('bytes32')) && !values[i].startsWith('CircuitValue256')) return false;
+        if (!(abiComponents[i].type.startsWith("uint") || abiComponents[i].type.startsWith("address") || abiComponents[i].type.startsWith("bytes") || abiComponents[i].type.startsWith("bool"))) return false;
+    }
+    return true;
+}
+
 export const getInputs = (inputs: string, inputSchema: string): any => {
     const inputSchemaJson = JSON.parse(inputSchema);
     const keys = Object.keys(inputSchemaJson);
-    const abi = getAbi();
+    const abis = getAbis().filter((x: any) => validateAbi(x, inputSchema));
+
+    if (abis.length === 0) {
+        throw new Error("Could not find valid ABI");
+    }
+
+    const abi = abis[0];
 
     const rawInputs: any = decodeAbiParameters(abi, inputs as `0x${string}`)[0];
+    const abiComponents = abi[0].components;
 
     const circuitInputs: any = {};
     for (let i = 0; i < keys.length; i++) {
         // if (keys[i] !== abi[i].name) throw new Error(`Input key ${keys[i]} does not match ABI name ${abi[i].name}`);
         if (Array.isArray(rawInputs[keys[i]])) {
-            circuitInputs[keys[i]] = rawInputs[keys[i]].map((x: any) => x.toString());
+            circuitInputs[keys[i]] = rawInputs[abiComponents[i].name].map((x: any) => x.toString());
         } else {
-            circuitInputs[keys[i]] = rawInputs[keys[i]].toString();
+            circuitInputs[keys[i]] = rawInputs[abiComponents[i].name].toString();
         }
     }
     return circuitInputs;

diff --git a/src/AxiomCli.sol b/src/AxiomCli.sol
@@ -4,5 +4,5 @@ pragma solidity ^0.8.0;
 library AxiomCli {
 
     /// @dev The SHA256 hash of the Axiom CLI binary
-    bytes public constant CLI_SHASUM = hex"27c283de43c8d73b186776da5374f57964a053df0ae49b09c5de696a2d30d9c2";
+    bytes public constant CLI_SHASUM = hex"d59ed924a489a673e35834bfb130a00601290a2a6a45b680358cfae2347ea3d1";
 }
diff --git a/test/Array.t.sol b/test/Array.t.sol
@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.19;
+
+import "../src/AxiomTest.sol";
+
+contract ArrayTest is AxiomTest {
+    using Axiom for Query;
+
+    struct AxiomInput {
+        uint64[] blockNumbers;
+        uint256[] slots;
+        address _address;
+    }
+
+    AxiomInput public input;
+    bytes32 public querySchema;
+
+    function setUp() public {
+        _createSelectForkAndSetupAxiom("sepolia", 5_103_100);
+
+        uint64[] memory blockNumbers = new uint64[](3);
+        uint256[] memory slots = new uint256[](3);
+        for (uint256 i = 0; i < 3; i++) {
+            blockNumbers[i] = 4_205_938;
+            slots[i] = i;
+        }
+        input = AxiomInput({ blockNumbers: blockNumbers, slots: slots, _address: address(0x8018fe32fCFd3d166E8b4c4E37105318A84BA11b)});
+
+        querySchema = axiomVm.readCircuit("test/circuit/array.circuit.ts");
+    }
+
+    /// @dev Simple demonstration of testing an Axiom client contract using Axiom cheatcodes
+    function test_simple_example() public {
+        // create a query into Axiom with default parameters
+        Query memory q = query(querySchema, abi.encode(input), address(0x8018fe32fCFd3d166E8b4c4E37105318A84BA11b));
+
+        // send the query to Axiom
+        q.send();
+    }
+}
diff --git a/test/circuit/array.circuit.ts b/test/circuit/array.circuit.ts
@@ -0,0 +1,32 @@
+import {
+    add,
+    sub,
+    mul,
+    div,
+    checkLessThan,
+    addToCallback,
+    CircuitValue,
+    CircuitValue256,
+    constant,
+    witness,
+    getAccount,
+    sum
+} from "@axiom-crypto/client";
+
+export interface CircuitInputs {
+    blockNumbers: CircuitValue[];
+    slots: CircuitValue256[];
+    address: CircuitValue;
+}
+
+/// Default inputs used for compilation.
+export const defaultInputs = {
+    "blockNumbers": [4205938, 4205938, 4205938],
+    "slots": [1, 2, 3],
+    "address": "0x8018fe32fCFd3d166E8b4c4E37105318A84BA11c"
+}
+
+export const circuit = async (inputs: CircuitInputs) => {
+    sum(inputs.blockNumbers);
+    addToCallback(inputs.address);
+};