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