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