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