@@ -32,12 +32,13 @@ thiserror = "1.0.26"
 color-eyre = "0.5"
 criterion = "0.3.4"
 
+cfg-if = "1.0"
+
 [dev-dependencies]
 hex-literal = "0.2.1"
 tokio = { version = "1.7.1", features = ["macros"] }
 serde_json = "1.0.64"
 ethers = { git = "https://github.com/gakonst/ethers-rs", features = ["abigen"] }
-cfg-if = "1.0"
 
 [[bench]]
 name = "groth16"
@@ -45,3 +46,4 @@ harness = false
 
 [features]
 bench-complex-all = []
+circom-2 = []

@@ -4,41 +4,92 @@
 #[derive(Clone, Debug)]
 pub struct Wasm(Instance);
 
-impl Wasm {
-    pub fn new(instance: Instance) -> Self {
-        Self(instance)
+pub trait CircomBase {
+    fn init(&self, sanity_check: bool) -> Result<()>;
+    fn func(&self, name: &str) -> &Function;
+    fn get_ptr_witness_buffer(&self) -> Result<i32>;
+    fn get_ptr_witness(&self, w: i32) -> Result<i32>;
+    fn get_n_vars(&self) -> Result<i32>;
+    fn get_signal_offset32(
+        &self,
+        p_sig_offset: u32,
+        component: u32,
+        hash_msb: u32,
+        hash_lsb: u32,
+    ) -> Result<()>;
+    fn set_signal(&self, c_idx: i32, component: i32, signal: i32, p_val: i32) -> Result<()>;
+    fn get_i32(&self, name: &str) -> Result<i32>;
+}
+
+pub trait Circom {
+    fn get_fr_len(&self) -> Result<i32>;
+    fn get_ptr_raw_prime(&self) -> Result<i32>;
+}
+
+pub trait Circom2 {
+    fn get_version(&self) -> Result<i32>;
+    fn get_field_num_len32(&self) -> Result<i32>;
+    fn get_raw_prime(&self) -> Result<()>;
+    fn read_shared_rw_memory(&self, i: i32) -> Result<i32>;
+}
+
+#[cfg(not(feature = "circom-2"))]
+impl Circom for Wasm {
+    fn get_fr_len(&self) -> Result<i32> {
+        self.get_i32("getFrLen")
+    }
+
+    fn get_ptr_raw_prime(&self) -> Result<i32> {
+        self.get_i32("getPRawPrime")
+    }
+}
+
+#[cfg(feature = "circom-2")]
+impl Circom2 for Wasm {
+    fn get_version(&self) -> Result<i32> {
+        self.get_i32("getVersion")
     }
 
-    pub fn init(&self, sanity_check: bool) -> Result<()> {
-        let func = self.func("init");
-        func.call(&[Value::I32(sanity_check as i32)])?;
-        Ok(())
+    fn get_field_num_len32(&self) -> Result<i32> {
+        self.get_i32("getFieldNumLen32")
     }
 
-    pub fn get_fr_len(&self) -> Result<i32> {
-        self.get_i32("getFrLen")
+    fn get_raw_prime(&self) -> Result<()> {
+        let func = self.func("getRawPrime");
+        let _result = func.call(&[])?;
+        Ok(())
     }
 
-    pub fn get_ptr_raw_prime(&self) -> Result<i32> {
-        self.get_i32("getPRawPrime")
+    fn read_shared_rw_memory(&self, i: i32) -> Result<i32> {
+        let func = self.func("readSharedRWMemory");
+        let result = func.call(&[i.into()])?;
+        Ok(result[0].unwrap_i32())
     }
+}
 
-    pub fn get_n_vars(&self) -> Result<i32> {
-        self.get_i32("getNVars")
+impl CircomBase for Wasm {
+    fn init(&self, sanity_check: bool) -> Result<()> {
+        let func = self.func("init");
+        func.call(&[Value::I32(sanity_check as i32)])?;
+        Ok(())
     }
 
-    pub fn get_ptr_witness_buffer(&self) -> Result<i32> {
+    fn get_ptr_witness_buffer(&self) -> Result<i32> {
         self.get_i32("getWitnessBuffer")
     }
 
-    pub fn get_ptr_witness(&self, w: i32) -> Result<i32> {
+    fn get_ptr_witness(&self, w: i32) -> Result<i32> {
         let func = self.func("getPWitness");
         let res = func.call(&[w.into()])?;
 
         Ok(res[0].unwrap_i32())
     }
 
-    pub fn get_signal_offset32(
+    fn get_n_vars(&self) -> Result<i32> {
+        self.get_i32("getNVars")
+    }
+
+    fn get_signal_offset32(
         &self,
         p_sig_offset: u32,
         component: u32,
@@ -56,7 +107,7 @@
         Ok(())
     }
 
-    pub fn set_signal(&self, c_idx: i32, component: i32, signal: i32, p_val: i32) -> Result<()> {
+    fn set_signal(&self, c_idx: i32, component: i32, signal: i32, p_val: i32) -> Result<()> {
         let func = self.func("setSignal");
         func.call(&[c_idx.into(), component.into(), signal.into(), p_val.into()])?;
 
@@ -76,3 +127,9 @@ .get_function(name)
             .unwrap_or_else(|_| panic!("function {} not found", name))
     }
 }
+
+impl Wasm {
+    pub fn new(instance: Instance) -> Self {
+        Self(instance)
+    }
+}

@@ -5,7 +5,13 @@ mod memory;
 pub(super) use memory::SafeMemory;
 
 mod circom;
-pub(super) use circom::Wasm;
+pub(super) use circom::{CircomBase, Wasm};
+
+#[cfg(feature = "circom-2")]
+pub(super) use circom::Circom2;
+
+#[cfg(not(feature = "circom-2"))]
+pub(super) use circom::Circom;
 
 use fnv::FnvHasher;
 use std::hash::Hasher;

@@ -1,10 +1,15 @@
+use super::{fnv, CircomBase, SafeMemory, Wasm};
 use color_eyre::Result;
 use num_bigint::BigInt;
 use num_traits::Zero;
 use std::cell::Cell;
 use wasmer::{imports, Function, Instance, Memory, MemoryType, Module, RuntimeError, Store};
 
-use super::{fnv, SafeMemory, Wasm};
+#[cfg(feature = "circom-2")]
+use super::Circom2;
+
+#[cfg(not(feature = "circom-2"))]
+use super::Circom;
 
 #[derive(Clone, Debug)]
 pub struct WitnessCalculator {
@@ -18,6 +23,16 @@ // From https://docs.wasmer.io/integrations/examples/exit-early
 #[derive(thiserror::Error, Debug, Clone, Copy)]
 #[error("{0}")]
 struct ExitCode(u32);
+
+#[cfg(feature = "circom-2")]
+fn from_array32(arr: Vec<i32>) -> BigInt {
+    let mut res = BigInt::zero();
+    let radix = BigInt::from(0x100000000u64);
+    for &val in arr.iter() {
+        res = res * &radix + BigInt::from(val);
+    }
+    res
+}
 
 impl WitnessCalculator {
     pub fn new(path: impl AsRef<std::path::Path>) -> Result<Self> {
@@ -38,22 +53,44 @@ "logGetSignal" => runtime::log_signal(&store),
                 "logFinishComponent" => runtime::log_component(&store),
                 "logStartComponent" => runtime::log_component(&store),
                 "log" => runtime::log_component(&store),
+                "exceptionHandler" => runtime::exception_handler(&store),
+                "showSharedRWMemory" => runtime::show_memory(&store),
             }
         };
         let instance = Wasm::new(Instance::new(&module, &import_object)?);
 
-        let n32 = (instance.get_fr_len()? >> 2) - 2;
+        let n32;
+        let prime: BigInt;
+        let mut safe_memory: SafeMemory;
 
-        let mut memory = SafeMemory::new(memory, n32 as usize, BigInt::zero());
+        cfg_if::cfg_if! {
+            if #[cfg(feature = "circom-2")] {
+                //let version = instance.get_version()?;
+                n32 = instance.get_field_num_len32()?;
+                safe_memory = SafeMemory::new(memory, n32 as usize, BigInt::zero());
+                let _res = instance.get_raw_prime()?;
+                let mut arr = vec![0; n32 as usize];
+                for i in 0..n32 {
+                    let res = instance.read_shared_rw_memory(i)?;
+                    arr[(n32 as usize) - (i as usize) - 1] = res;
+                }
+                prime = from_array32(arr);
+            } else {
+                // Fallback to Circom 1 behavior
+                //version = 1;
+                n32 = (instance.get_fr_len()? >> 2) - 2;
+                safe_memory = SafeMemory::new(memory, n32 as usize, BigInt::zero());
+                let ptr = instance.get_ptr_raw_prime()?;
+                prime = safe_memory.read_big(ptr as usize, n32 as usize)?;
+            }
+        }
 
-        let ptr = instance.get_ptr_raw_prime()?;
-        let prime = memory.read_big(ptr as usize, n32 as usize)?;
         let n64 = ((prime.bits() - 1) / 64 + 1) as i32;
-        memory.prime = prime;
+        safe_memory.prime = prime;
 
         Ok(WitnessCalculator {
             instance,
-            memory,
+            memory: safe_memory,
             n64,
         })
     }
@@ -159,6 +196,20 @@ a, b, c, d, e, f
             );
             RuntimeError::raise(Box::new(ExitCode(1)));
         }
+        Function::new_native(store, func)
+    }
+
+    // Circom 2.0
+    pub fn exception_handler(store: &Store) -> Function {
+        #[allow(unused)]
+        fn func(a: i32) {}
+        Function::new_native(store, func)
+    }
+
+    // Circom 2.0
+    pub fn show_memory(store: &Store) -> Function {
+        #[allow(unused)]
+        fn func() {}
         Function::new_native(store, func)
     }
 

@@ -58,3 +58,35 @@ let _params = generate_random_parameters::<Bn254, _, _>(circom, &mut rng).unwrap();
 
     builder.build().unwrap_err();
 }
+
+#[test]
+#[cfg(feature = "circom-2")]
+fn groth16_proof_circom2() -> Result<()> {
+    let cfg = CircomConfig::<Bn254>::new(
+        "./test-vectors/circom2_multiplier2.wasm",
+        "./test-vectors/circom2_multiplier2.r1cs",
+    )?;
+    let mut builder = CircomBuilder::new(cfg);
+    builder.push_input("a", 3);
+    builder.push_input("b", 11);
+
+    // create an empty instance for setting it up
+    let circom = builder.setup();
+
+    let mut rng = thread_rng();
+    let params = generate_random_parameters::<Bn254, _, _>(circom, &mut rng)?;
+
+    let circom = builder.build()?;
+
+    let inputs = circom.get_public_inputs().unwrap();
+
+    let proof = prove(circom, &params, &mut rng)?;
+
+    let pvk = prepare_verifying_key(&params.vk);
+
+    let verified = verify_proof(&pvk, &proof, &inputs)?;
+
+    assert!(verified);
+
+    Ok(())
+}