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