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