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