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