1use color_eyre::Result;
  2use wasmer::{Function, Instance, Value};
  3
  4#[derive(Clone, Debug)]
  5pub struct Wasm(Instance);
  6
  7pub trait CircomBase {
  8    fn init(&self, sanity_check: bool) -> Result<()>;
  9    fn func(&self, name: &str) -> &Function;
 10    fn get_ptr_witness_buffer(&self) -> Result<i32>;
 11    fn get_ptr_witness(&self, w: i32) -> Result<i32>;
 12    fn get_n_vars(&self) -> Result<i32>;
 13    fn get_signal_offset32(
 14        &self,
 15        p_sig_offset: u32,
 16        component: u32,
 17        hash_msb: u32,
 18        hash_lsb: u32,
 19    ) -> Result<()>;
 20    fn set_signal(&self, c_idx: i32, component: i32, signal: i32, p_val: i32) -> Result<()>;
 21    fn get_i32(&self, name: &str) -> Result<i32>;
 22    // Only exists natively in Circom2, hardcoded for Circom
 23    fn get_version(&self) -> Result<i32>;
 24}
 25
 26pub trait Circom {
 27    fn get_fr_len(&self) -> Result<i32>;
 28    fn get_ptr_raw_prime(&self) -> Result<i32>;
 29}
 30
 31pub trait Circom2 {
 32    fn get_field_num_len32(&self) -> Result<i32>;
 33    fn get_raw_prime(&self) -> Result<()>;
 34    fn read_shared_rw_memory(&self, i: i32) -> Result<i32>;
 35    fn write_shared_rw_memory(&self, i: i32, v: i32) -> Result<()>;
 36    fn set_input_signal(&self, hmsb: i32, hlsb: i32, pos: i32) -> Result<()>;
 37    fn get_witness(&self, i: i32) -> Result<()>;
 38    fn get_witness_size(&self) -> Result<i32>;
 39}
 40
 41impl Circom for Wasm {
 42    fn get_fr_len(&self) -> Result<i32> {
 43        self.get_i32("getFrLen")
 44    }
 45
 46    fn get_ptr_raw_prime(&self) -> Result<i32> {
 47        self.get_i32("getPRawPrime")
 48    }
 49}
 50
 51#[cfg(feature = "circom-2")]
 52impl Circom2 for Wasm {
 53    fn get_field_num_len32(&self) -> Result<i32> {
 54        self.get_i32("getFieldNumLen32")
 55    }
 56
 57    fn get_raw_prime(&self) -> Result<()> {
 58        let func = self.func("getRawPrime");
 59        func.call(&[])?;
 60        Ok(())
 61    }
 62
 63    fn read_shared_rw_memory(&self, i: i32) -> Result<i32> {
 64        let func = self.func("readSharedRWMemory");
 65        let result = func.call(&[i.into()])?;
 66        Ok(result[0].unwrap_i32())
 67    }
 68
 69    fn write_shared_rw_memory(&self, i: i32, v: i32) -> Result<()> {
 70        let func = self.func("writeSharedRWMemory");
 71        func.call(&[i.into(), v.into()])?;
 72        Ok(())
 73    }
 74
 75    fn set_input_signal(&self, hmsb: i32, hlsb: i32, pos: i32) -> Result<()> {
 76        let func = self.func("setInputSignal");
 77        func.call(&[hmsb.into(), hlsb.into(), pos.into()])?;
 78        Ok(())
 79    }
 80
 81    fn get_witness(&self, i: i32) -> Result<()> {
 82        let func = self.func("getWitness");
 83        func.call(&[i.into()])?;
 84        Ok(())
 85    }
 86
 87    fn get_witness_size(&self) -> Result<i32> {
 88        let func = self.func("getWitnessSize");
 89        let result = func.call(&[])?;
 90        Ok(result[0].unwrap_i32())
 91    }
 92}
 93
 94impl CircomBase for Wasm {
 95    fn init(&self, sanity_check: bool) -> Result<()> {
 96        let func = self.func("init");
 97        func.call(&[Value::I32(sanity_check as i32)])?;
 98        Ok(())
 99    }
100
101    fn get_ptr_witness_buffer(&self) -> Result<i32> {
102        self.get_i32("getWitnessBuffer")
103    }
104
105    fn get_ptr_witness(&self, w: i32) -> Result<i32> {
106        let func = self.func("getPWitness");
107        let res = func.call(&[w.into()])?;
108
109        Ok(res[0].unwrap_i32())
110    }
111
112    fn get_n_vars(&self) -> Result<i32> {
113        self.get_i32("getNVars")
114    }
115
116    fn get_signal_offset32(
117        &self,
118        p_sig_offset: u32,
119        component: u32,
120        hash_msb: u32,
121        hash_lsb: u32,
122    ) -> Result<()> {
123        let func = self.func("getSignalOffset32");
124        func.call(&[
125            p_sig_offset.into(),
126            component.into(),
127            hash_msb.into(),
128            hash_lsb.into(),
129        ])?;
130
131        Ok(())
132    }
133
134    fn set_signal(&self, c_idx: i32, component: i32, signal: i32, p_val: i32) -> Result<()> {
135        let func = self.func("setSignal");
136        func.call(&[c_idx.into(), component.into(), signal.into(), p_val.into()])?;
137
138        Ok(())
139    }
140
141    // Default to version 1 if it isn't explicitly defined
142    fn get_version(&self) -> Result<i32> {
143        match self.0.exports.get_function("getVersion") {
144            Ok(func) => Ok(func.call(&[])?[0].unwrap_i32()),
145            Err(_) => Ok(1),
146        }
147    }
148
149    fn get_i32(&self, name: &str) -> Result<i32> {
150        let func = self.func(name);
151        let result = func.call(&[])?;
152        Ok(result[0].unwrap_i32())
153    }
154
155    fn func(&self, name: &str) -> &Function {
156        self.0
157            .exports
158            .get_function(name)
159            .unwrap_or_else(|_| panic!("function {} not found", name))
160    }
161}
162
163impl Wasm {
164    pub fn new(instance: Instance) -> Self {
165        Self(instance)
166    }
167}