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