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