From dbbb43aec8605ea71b98c9ad8088adc9751d5b80 Mon Sep 17 00:00:00 2001
From: Hrik Bhowal <23041116+HrikB@users.noreply.github.com>
Date: Fri, 7 Jun 2024 17:36:05 -0700
Subject: [PATCH] fix: parameterize circuit paths and fix tests

---
 src/AxiomVm.sol       | 52 ++++++++++++++++++++++++++++++-------------
 test/AxiomStdRs.t.sol |  7 +++++-
 2 files changed, 43 insertions(+), 16 deletions(-)

diff --git a/src/AxiomVm.sol b/src/AxiomVm.sol
index 8519f47..48a3c8f 100644
--- a/src/AxiomVm.sol
+++ b/src/AxiomVm.sol
@@ -139,9 +139,18 @@ contract AxiomVm is Test {
     /// @dev Default dummy query schema for Rust circuits
     bytes32 constant DEFAULT_RUST_QUERY_SCHEMA = 0xdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeefdeadbeef;
 
+    /// @dev Path to the manifest file
+    string public manifestPath;
+
     /// @dev Circuit path for Rust circuits
     string public rustCircuitPath;
 
+    /// @dev Path to the data directory
+    string public dataPath;
+
+    /// @dev Binary name for Rust circuits
+    string public bin;
+
     constructor(address _axiomV2QueryAddress, string memory _urlOrAlias) {
         axiomV2QueryAddress = _axiomV2QueryAddress;
         urlOrAlias = _urlOrAlias;
@@ -267,9 +276,19 @@ contract AxiomVm is Test {
      * @param _rustCircuitPath path to the Rust circuit file
      * @return querySchema
      */
-    function readRustCircuit(string memory _rustCircuitPath) public returns (bytes32 querySchema) {
+    function readRustCircuit(
+        string memory _manifestPath,
+        string memory _rustCircuitPath,
+        string memory _dataPath,
+        string memory _bin
+    ) public returns (bytes32 querySchema) {
         isRustCircuit = true;
+
+        manifestPath = _manifestPath;
         rustCircuitPath = _rustCircuitPath;
+        dataPath = _dataPath;
+        bin = _bin;
+
         querySchema = DEFAULT_RUST_QUERY_SCHEMA;
     }
 
@@ -554,23 +573,26 @@ contract AxiomVm is Test {
         IAxiomV2Query.AxiomV2FeeData memory feeData
     ) internal returns (string memory output) {
         // Get compute results from Rust circuit
-        string[] memory witnessGen = new string[](16);
+        string[] memory witnessGen = new string[](18);
+
         witnessGen[0] = "cargo";
         witnessGen[1] = "+nightly-2024-01-01";
         witnessGen[2] = "run";
-        witnessGen[3] = "--manifest-path";
-        witnessGen[4] = "test/circuit-rs/Cargo.toml";
-        witnessGen[5] = "--";
-        witnessGen[6] = "--input";
-        witnessGen[7] = "test/circuit-rs/data/account_age_input.json";
-        witnessGen[8] = "--data-path";
-        witnessGen[9] = "test/circuit-rs/data";
-        witnessGen[10] = "-k";
-        witnessGen[11] = "12";
-        witnessGen[12] = "-p";
-        witnessGen[13] = vm.rpcUrl(urlOrAlias);
-        witnessGen[14] = "--to-stdout";
-        witnessGen[15] = "witness-gen";
+        witnessGen[3] = "--bin";
+        witnessGen[4] = bin;
+        witnessGen[5] = "--manifest-path";
+        witnessGen[6] = manifestPath; // "test/circuit-rs/Cargo.toml"
+        witnessGen[7] = "--";
+        witnessGen[8] = "--input";
+        witnessGen[9] = rustCircuitPath; // "test/circuit-rs/data/account_age_input.json"
+        witnessGen[10] = "--data-path";
+        witnessGen[11] = dataPath; // "test/circuit-rs/data"
+        witnessGen[12] = "-k";
+        witnessGen[13] = "12";
+        witnessGen[14] = "-p";
+        witnessGen[15] = vm.rpcUrl(urlOrAlias);
+        witnessGen[16] = "--to-stdout";
+        witnessGen[17] = "witness-gen";
 
         bytes memory axiomOutput0 = vm.ffi(witnessGen);
         string memory computeResults = string(axiomOutput0);
diff --git a/test/AxiomStdRs.t.sol b/test/AxiomStdRs.t.sol
index fa035fe..d8c4c7b 100644
--- a/test/AxiomStdRs.t.sol
+++ b/test/AxiomStdRs.t.sol
@@ -26,7 +26,12 @@ contract AxiomStdRsTest is AxiomTest {
 
         defaultInput =
             Input({ blockNumber: 5_146_659, _address: uint256(uint160(0xef663bB0e0b1091571DaD9715994bc81e9f5a2ab)) });
-        querySchema = axiomVm.readRustCircuit("test/circuit-rs/account_age.rs");
+        querySchema = axiomVm.readRustCircuit({
+            _manifestPath: "test/circuit-rs/Cargo.toml",
+            _rustCircuitPath: "test/circuit-rs/data/account_age_input.json",
+            _dataPath: "test/circuit-rs/data",
+            _bin: "circuit-rs"
+        });
         accountAge = new AccountAge(axiomV2QueryAddress, uint64(block.chainid), querySchema);
 
         callbackExtraData = bytes("");