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}
 23
 24pub trait Circom {
 25    fn get_fr_len(&self) -> Result<i32>;
 26    fn get_ptr_raw_prime(&self) -> Result<i32>;
 27}
 28
 29pub trait Circom2 {
 30    fn get_version(&self) -> Result<i32>;
 31    fn get_field_num_len32(&self) -> Result<i32>;
 32    fn get_raw_prime(&self) -> Result<()>;
 33    fn read_shared_rw_memory(&self, i: i32) -> Result<i32>;
 34    fn write_shared_rw_memory(&self, i: i32, v: i32) -> Result<()>;
 35    fn set_input_signal(&self, hmsb: i32, hlsb: i32, pos: i32) -> Result<()>;
 36    fn get_witness(&self, i: i32) -> Result<()>;
 37    fn get_witness_size(&self) -> Result<i32>;
 38}
 39
 40#[cfg(not(feature = "circom-2"))]
 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_version(&self) -> Result<i32> {
 54        self.get_i32("getVersion")
 55    }
 56
 57    fn get_field_num_len32(&self) -> Result<i32> {
 58        self.get_i32("getFieldNumLen32")
 59    }
 60
 61    fn get_raw_prime(&self) -> Result<()> {
 62        let func = self.func("getRawPrime");
 63        func.call(&[])?;
 64        Ok(())
 65    }
 66
 67    fn read_shared_rw_memory(&self, i: i32) -> Result<i32> {
 68        let func = self.func("readSharedRWMemory");
 69        let result = func.call(&[i.into()])?;
 70        Ok(result[0].unwrap_i32())
 71    }
 72
 73    fn write_shared_rw_memory(&self, i: i32, v: i32) -> Result<()> {
 74        let func = self.func("writeSharedRWMemory");
 75        func.call(&[i.into(), v.into()])?;
 76        Ok(())
 77    }
 78
 79    fn set_input_signal(&self, hmsb: i32, hlsb: i32, pos: i32) -> Result<()> {
 80        let func = self.func("setInputSignal");
 81        func.call(&[hmsb.into(), hlsb.into(), pos.into()])?;
 82        Ok(())
 83    }
 84
 85    fn get_witness(&self, i: i32) -> Result<()> {
 86        let func = self.func("getWitness");
 87        func.call(&[i.into()])?;
 88        Ok(())
 89    }
 90
 91    fn get_witness_size(&self) -> Result<i32> {
 92        let func = self.func("getWitnessSize");
 93        let result = func.call(&[])?;
 94        Ok(result[0].unwrap_i32())
 95    }
 96}
 97
 98impl CircomBase for Wasm {
 99    fn init(&self, sanity_check: bool) -> Result<()> {
100        let func = self.func("init");
101        func.call(&[Value::I32(sanity_check as i32)])?;
102        Ok(())
103    }
104
105    fn get_ptr_witness_buffer(&self) -> Result<i32> {
106        self.get_i32("getWitnessBuffer")
107    }
108
109    fn get_ptr_witness(&self, w: i32) -> Result<i32> {
110        let func = self.func("getPWitness");
111        let res = func.call(&[w.into()])?;
112
113        Ok(res[0].unwrap_i32())
114    }
115
116    fn get_n_vars(&self) -> Result<i32> {
117        self.get_i32("getNVars")
118    }
119
120    fn get_signal_offset32(
121        &self,
122        p_sig_offset: u32,
123        component: u32,
124        hash_msb: u32,
125        hash_lsb: u32,
126    ) -> Result<()> {
127        let func = self.func("getSignalOffset32");
128        func.call(&[
129            p_sig_offset.into(),
130            component.into(),
131            hash_msb.into(),
132            hash_lsb.into(),
133        ])?;
134
135        Ok(())
136    }
137
138    fn set_signal(&self, c_idx: i32, component: i32, signal: i32, p_val: i32) -> Result<()> {
139        let func = self.func("setSignal");
140        func.call(&[c_idx.into(), component.into(), signal.into(), p_val.into()])?;
141
142        Ok(())
143    }
144
145    fn get_i32(&self, name: &str) -> Result<i32> {
146        let func = self.func(name);
147        let result = func.call(&[])?;
148        Ok(result[0].unwrap_i32())
149    }
150
151    fn func(&self, name: &str) -> &Function {
152        self.0
153            .exports
154            .get_function(name)
155            .unwrap_or_else(|_| panic!("function {} not found", name))
156    }
157}
158
159impl Wasm {
160    pub fn new(instance: Instance) -> Self {
161        Self(instance)
162    }
163}