1use super::{fnv, CircomBase, SafeMemory, Wasm};
  2use color_eyre::Result;
  3use num_bigint::BigInt;
  4use num_traits::Zero;
  5use std::cell::Cell;
  6use wasmer::{imports, Function, Instance, Memory, MemoryType, Module, RuntimeError, Store};
  7
  8#[cfg(feature = "circom-2")]
  9use super::Circom2;
 10
 11#[cfg(not(feature = "circom-2"))]
 12use super::Circom;
 13
 14#[derive(Clone, Debug)]
 15pub struct WitnessCalculator {
 16    pub instance: Wasm,
 17    pub memory: SafeMemory,
 18    pub n64: i32,
 19}
 20
 21// Error type to signal end of execution.
 22// From https://docs.wasmer.io/integrations/examples/exit-early
 23#[derive(thiserror::Error, Debug, Clone, Copy)]
 24#[error("{0}")]
 25struct ExitCode(u32);
 26
 27#[cfg(feature = "circom-2")]
 28fn from_array32(arr: Vec<i32>) -> BigInt {
 29    let mut res = BigInt::zero();
 30    let radix = BigInt::from(0x100000000u64);
 31    for &val in arr.iter() {
 32        res = res * &radix + BigInt::from(val);
 33    }
 34    res
 35}
 36
 37impl WitnessCalculator {
 38    pub fn new(path: impl AsRef<std::path::Path>) -> Result<Self> {
 39        let store = Store::default();
 40        let module = Module::from_file(&store, path)?;
 41
 42        // Set up the memory
 43        let memory = Memory::new(&store, MemoryType::new(2000, None, false)).unwrap();
 44        let import_object = imports! {
 45            "env" => {
 46                "memory" => memory.clone(),
 47            },
 48            // Host function callbacks from the WASM
 49            "runtime" => {
 50                "error" => runtime::error(&store),
 51                "logSetSignal" => runtime::log_signal(&store),
 52                "logGetSignal" => runtime::log_signal(&store),
 53                "logFinishComponent" => runtime::log_component(&store),
 54                "logStartComponent" => runtime::log_component(&store),
 55                "log" => runtime::log_component(&store),
 56                "exceptionHandler" => runtime::exception_handler(&store),
 57                "showSharedRWMemory" => runtime::show_memory(&store),
 58            }
 59        };
 60        let instance = Wasm::new(Instance::new(&module, &import_object)?);
 61
 62        let n32;
 63        let prime: BigInt;
 64        let mut safe_memory: SafeMemory;
 65
 66        cfg_if::cfg_if! {
 67            if #[cfg(feature = "circom-2")] {
 68                //let version = instance.get_version()?;
 69                n32 = instance.get_field_num_len32()?;
 70                safe_memory = SafeMemory::new(memory, n32 as usize, BigInt::zero());
 71                let _res = instance.get_raw_prime()?;
 72                let mut arr = vec![0; n32 as usize];
 73                for i in 0..n32 {
 74                    let res = instance.read_shared_rw_memory(i)?;
 75                    arr[(n32 as usize) - (i as usize) - 1] = res;
 76                }
 77                prime = from_array32(arr);
 78            } else {
 79                // Fallback to Circom 1 behavior
 80                //version = 1;
 81                n32 = (instance.get_fr_len()? >> 2) - 2;
 82                safe_memory = SafeMemory::new(memory, n32 as usize, BigInt::zero());
 83                let ptr = instance.get_ptr_raw_prime()?;
 84                prime = safe_memory.read_big(ptr as usize, n32 as usize)?;
 85            }
 86        }
 87
 88        let n64 = ((prime.bits() - 1) / 64 + 1) as i32;
 89        safe_memory.prime = prime;
 90
 91        Ok(WitnessCalculator {
 92            instance,
 93            memory: safe_memory,
 94            n64,
 95        })
 96    }
 97
 98    pub fn calculate_witness<I: IntoIterator<Item = (String, Vec<BigInt>)>>(
 99        &mut self,
100        inputs: I,
101        sanity_check: bool,
102    ) -> Result<Vec<BigInt>> {
103        let old_mem_free_pos = self.memory.free_pos();
104
105        self.instance.init(sanity_check)?;
106
107        let p_sig_offset = self.memory.alloc_u32();
108        let p_fr = self.memory.alloc_fr();
109
110        // allocate the inputs
111        for (name, values) in inputs.into_iter() {
112            let (msb, lsb) = fnv(&name);
113            self.instance
114                .get_signal_offset32(p_sig_offset, 0, msb, lsb)?;
115
116            let sig_offset = self.memory.read_u32(p_sig_offset as usize) as usize;
117
118            for (i, value) in values.into_iter().enumerate() {
119                self.memory.write_fr(p_fr as usize, &value)?;
120                self.instance
121                    .set_signal(0, 0, (sig_offset + i) as i32, p_fr as i32)?;
122            }
123        }
124
125        let mut w = Vec::new();
126        let n_vars = self.instance.get_n_vars()?;
127        for i in 0..n_vars {
128            let ptr = self.instance.get_ptr_witness(i)? as usize;
129            let el = self.memory.read_fr(ptr)?;
130            w.push(el);
131        }
132
133        self.memory.set_free_pos(old_mem_free_pos);
134
135        Ok(w)
136    }
137
138    pub fn calculate_witness_element<
139        E: ark_ec::PairingEngine,
140        I: IntoIterator<Item = (String, Vec<BigInt>)>,
141    >(
142        &mut self,
143        inputs: I,
144        sanity_check: bool,
145    ) -> Result<Vec<E::Fr>> {
146        use ark_ff::{FpParameters, PrimeField};
147        let witness = self.calculate_witness(inputs, sanity_check)?;
148        let modulus = <<E::Fr as PrimeField>::Params as FpParameters>::MODULUS;
149
150        // convert it to field elements
151        use num_traits::Signed;
152        let witness = witness
153            .into_iter()
154            .map(|w| {
155                let w = if w.sign() == num_bigint::Sign::Minus {
156                    // Need to negate the witness element if negative
157                    modulus.into() - w.abs().to_biguint().unwrap()
158                } else {
159                    w.to_biguint().unwrap()
160                };
161                E::Fr::from(w)
162            })
163            .collect::<Vec<_>>();
164
165        Ok(witness)
166    }
167
168    pub fn get_witness_buffer(&self) -> Result<Vec<u8>> {
169        let ptr = self.instance.get_ptr_witness_buffer()? as usize;
170
171        let view = self.memory.memory.view::<u8>();
172
173        let len = self.instance.get_n_vars()? * self.n64 * 8;
174        let arr = view[ptr..ptr + len as usize]
175            .iter()
176            .map(Cell::get)
177            .collect::<Vec<_>>();
178
179        Ok(arr)
180    }
181}
182
183// callback hooks for debugging
184mod runtime {
185    use super::*;
186
187    pub fn error(store: &Store) -> Function {
188        #[allow(unused)]
189        #[allow(clippy::many_single_char_names)]
190        fn func(a: i32, b: i32, c: i32, d: i32, e: i32, f: i32) {
191            // NOTE: We can also get more information why it is failing, see p2str etc here:
192            // https://github.com/iden3/circom_runtime/blob/master/js/witness_calculator.js#L52-L64
193            println!(
194                "runtime error, exiting early: {0} {1} {2} {3} {4} {5}",
195                a, b, c, d, e, f
196            );
197            RuntimeError::raise(Box::new(ExitCode(1)));
198        }
199        Function::new_native(store, func)
200    }
201
202    // Circom 2.0
203    pub fn exception_handler(store: &Store) -> Function {
204        #[allow(unused)]
205        fn func(a: i32) {}
206        Function::new_native(store, func)
207    }
208
209    // Circom 2.0
210    pub fn show_memory(store: &Store) -> Function {
211        #[allow(unused)]
212        fn func() {}
213        Function::new_native(store, func)
214    }
215
216    pub fn log_signal(store: &Store) -> Function {
217        #[allow(unused)]
218        fn func(a: i32, b: i32) {}
219        Function::new_native(store, func)
220    }
221
222    pub fn log_component(store: &Store) -> Function {
223        #[allow(unused)]
224        fn func(a: i32) {}
225        Function::new_native(store, func)
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use std::{collections::HashMap, path::PathBuf};
233
234    struct TestCase<'a> {
235        circuit_path: &'a str,
236        inputs_path: &'a str,
237        n_vars: u32,
238        n64: u32,
239        witness: &'a [&'a str],
240    }
241
242    pub fn root_path(p: &str) -> String {
243        let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
244        path.push(p);
245        path.to_string_lossy().to_string()
246    }
247
248    #[test]
249    fn multiplier_1() {
250        run_test(TestCase {
251            circuit_path: root_path("test-vectors/mycircuit.wasm").as_str(),
252            inputs_path: root_path("test-vectors/mycircuit-input1.json").as_str(),
253            n_vars: 4,
254            n64: 4,
255            witness: &["1", "33", "3", "11"],
256        });
257    }
258
259    #[test]
260    fn multiplier_2() {
261        run_test(TestCase {
262            circuit_path: root_path("test-vectors/mycircuit.wasm").as_str(),
263            inputs_path: root_path("test-vectors/mycircuit-input2.json").as_str(),
264            n_vars: 4,
265            n64: 4,
266            witness: &[
267                "1",
268                "21888242871839275222246405745257275088548364400416034343698204186575672693159",
269                "21888242871839275222246405745257275088548364400416034343698204186575796149939",
270                "11",
271            ],
272        });
273    }
274
275    #[test]
276    fn multiplier_3() {
277        run_test(TestCase {
278            circuit_path: root_path("test-vectors/mycircuit.wasm").as_str(),
279            inputs_path: root_path("test-vectors/mycircuit-input3.json").as_str(),
280            n_vars: 4,
281            n64: 4,
282            witness: &[
283                "1",
284                "21888242871839275222246405745257275088548364400416034343698204186575808493616",
285                "10944121435919637611123202872628637544274182200208017171849102093287904246808",
286                "2",
287            ],
288        });
289    }
290
291    #[test]
292    fn safe_multipler() {
293        let witness =
294            std::fs::read_to_string(&root_path("test-vectors/safe-circuit-witness.json")).unwrap();
295        let witness: Vec<String> = serde_json::from_str(&witness).unwrap();
296        let witness = &witness.iter().map(|x| x.as_ref()).collect::<Vec<_>>();
297        run_test(TestCase {
298            circuit_path: root_path("test-vectors/circuit2.wasm").as_str(),
299            inputs_path: root_path("test-vectors/mycircuit-input1.json").as_str(),
300            n_vars: 132, // 128 + 4
301            n64: 4,
302            witness,
303        });
304    }
305
306    #[test]
307    fn smt_verifier() {
308        let witness =
309            std::fs::read_to_string(&root_path("test-vectors/smtverifier10-witness.json")).unwrap();
310        let witness: Vec<String> = serde_json::from_str(&witness).unwrap();
311        let witness = &witness.iter().map(|x| x.as_ref()).collect::<Vec<_>>();
312
313        run_test(TestCase {
314            circuit_path: root_path("test-vectors/smtverifier10.wasm").as_str(),
315            inputs_path: root_path("test-vectors/smtverifier10-input.json").as_str(),
316            n_vars: 4794,
317            n64: 4,
318            witness,
319        });
320    }
321
322    use serde_json::Value;
323    use std::str::FromStr;
324
325    fn value_to_bigint(v: Value) -> BigInt {
326        match v {
327            Value::String(inner) => BigInt::from_str(&inner).unwrap(),
328            Value::Number(inner) => BigInt::from(inner.as_u64().expect("not a u32")),
329            _ => panic!("unsupported type"),
330        }
331    }
332
333    fn run_test(case: TestCase) {
334        let mut wtns = WitnessCalculator::new(case.circuit_path).unwrap();
335        assert_eq!(
336            wtns.memory.prime.to_str_radix(16),
337            "30644E72E131A029B85045B68181585D2833E84879B9709143E1F593F0000001".to_lowercase()
338        );
339        assert_eq!(wtns.instance.get_n_vars().unwrap() as u32, case.n_vars);
340        assert_eq!(wtns.n64 as u32, case.n64);
341
342        let inputs_str = std::fs::read_to_string(case.inputs_path).unwrap();
343        let inputs: std::collections::HashMap<String, serde_json::Value> =
344            serde_json::from_str(&inputs_str).unwrap();
345
346        let inputs = inputs
347            .iter()
348            .map(|(key, value)| {
349                let res = match value {
350                    Value::String(inner) => {
351                        vec![BigInt::from_str(inner).unwrap()]
352                    }
353                    Value::Number(inner) => {
354                        vec![BigInt::from(inner.as_u64().expect("not a u32"))]
355                    }
356                    Value::Array(inner) => inner.iter().cloned().map(value_to_bigint).collect(),
357                    _ => panic!(),
358                };
359
360                (key.clone(), res)
361            })
362            .collect::<HashMap<_, _>>();
363
364        let res = wtns.calculate_witness(inputs, false).unwrap();
365        for (r, w) in res.iter().zip(case.witness) {
366            assert_eq!(r, &BigInt::from_str(w).unwrap());
367        }
368    }
369}