1use color_eyre::Result;
  2use num_bigint::BigInt;
  3use num_traits::Zero;
  4use std::cell::Cell;
  5use wasmer::{imports, Function, Instance, Memory, MemoryType, Module, Store};
  6
  7use super::{fnv, SafeMemory, Wasm};
  8
  9#[derive(Clone, Debug)]
 10pub struct WitnessCalculator {
 11    pub instance: Wasm,
 12    pub memory: SafeMemory,
 13    pub n64: i32,
 14}
 15
 16impl WitnessCalculator {
 17    pub fn new(path: impl AsRef<std::path::Path>) -> Result<Self> {
 18        let store = Store::default();
 19        let module = Module::from_file(&store, path)?;
 20
 21        // Set up the memory
 22        let memory = Memory::new(&store, MemoryType::new(2000, None, false)).unwrap();
 23        let import_object = imports! {
 24            "env" => {
 25                "memory" => memory.clone(),
 26            },
 27            // Host function callbacks from the WASM
 28            "runtime" => {
 29                "error" => runtime::error(&store),
 30                "logSetSignal" => runtime::log_signal(&store),
 31                "logGetSignal" => runtime::log_signal(&store),
 32                "logFinishComponent" => runtime::log_component(&store),
 33                "logStartComponent" => runtime::log_component(&store),
 34                "log" => runtime::log_component(&store),
 35            }
 36        };
 37        let instance = Wasm::new(Instance::new(&module, &import_object)?);
 38
 39        let n32 = (instance.get_fr_len()? >> 2) - 2;
 40
 41        let mut memory = SafeMemory::new(memory, n32 as usize, BigInt::zero());
 42
 43        let ptr = instance.get_ptr_raw_prime()?;
 44        let prime = memory.read_big(ptr as usize, n32 as usize)?;
 45        let n64 = ((prime.bits() - 1) / 64 + 1) as i32;
 46        memory.prime = prime;
 47
 48        Ok(WitnessCalculator {
 49            instance,
 50            memory,
 51            n64,
 52        })
 53    }
 54
 55    pub fn calculate_witness<I: IntoIterator<Item = (String, Vec<BigInt>)>>(
 56        &mut self,
 57        inputs: I,
 58        sanity_check: bool,
 59    ) -> Result<Vec<BigInt>> {
 60        let old_mem_free_pos = self.memory.free_pos();
 61
 62        self.instance.init(sanity_check)?;
 63
 64        let p_sig_offset = self.memory.alloc_u32();
 65        let p_fr = self.memory.alloc_fr();
 66
 67        // allocate the inputs
 68        for (name, values) in inputs.into_iter() {
 69            let (msb, lsb) = fnv(&name);
 70            self.instance
 71                .get_signal_offset32(p_sig_offset, 0, msb, lsb)?;
 72
 73            let sig_offset = self.memory.read_u32(p_sig_offset as usize) as usize;
 74
 75            for (i, value) in values.into_iter().enumerate() {
 76                self.memory.write_fr(p_fr as usize, &value)?;
 77                self.instance
 78                    .set_signal(0, 0, (sig_offset + i) as i32, p_fr as i32)?;
 79            }
 80        }
 81
 82        let mut w = Vec::new();
 83        let n_vars = self.instance.get_n_vars()?;
 84        for i in 0..n_vars {
 85            let ptr = self.instance.get_ptr_witness(i)? as usize;
 86            let el = self.memory.read_fr(ptr)?;
 87            w.push(el);
 88        }
 89
 90        self.memory.set_free_pos(old_mem_free_pos);
 91
 92        Ok(w)
 93    }
 94
 95    pub fn calculate_witness_element<
 96        E: ark_ec::PairingEngine,
 97        I: IntoIterator<Item = (String, Vec<BigInt>)>,
 98    >(
 99        &mut self,
100        inputs: I,
101        sanity_check: bool,
102    ) -> Result<Vec<E::Fr>> {
103        use ark_ff::{FpParameters, PrimeField};
104        let witness = self.calculate_witness(inputs, sanity_check)?;
105        let modulus = <<E::Fr as PrimeField>::Params as FpParameters>::MODULUS;
106
107        // convert it to field elements
108        use num_traits::Signed;
109        let witness = witness
110            .into_iter()
111            .map(|w| {
112                let w = if w.sign() == num_bigint::Sign::Minus {
113                    // Need to negate the witness element if negative
114                    modulus.into() - w.abs().to_biguint().unwrap()
115                } else {
116                    w.to_biguint().unwrap()
117                };
118                E::Fr::from(w)
119            })
120            .collect::<Vec<_>>();
121
122        Ok(witness)
123    }
124
125    pub fn get_witness_buffer(&self) -> Result<Vec<u8>> {
126        let ptr = self.instance.get_ptr_witness_buffer()? as usize;
127
128        let view = self.memory.memory.view::<u8>();
129
130        let len = self.instance.get_n_vars()? * self.n64 * 8;
131        let arr = view[ptr..ptr + len as usize]
132            .iter()
133            .map(Cell::get)
134            .collect::<Vec<_>>();
135
136        Ok(arr)
137    }
138}
139
140// callback hooks for debugging
141mod runtime {
142    use super::*;
143
144    pub fn error(store: &Store) -> Function {
145        #[allow(unused)]
146        #[allow(clippy::many_single_char_names)]
147        fn func(a: i32, b: i32, c: i32, d: i32, e: i32, f: i32) {}
148        Function::new_native(store, func)
149    }
150
151    pub fn log_signal(store: &Store) -> Function {
152        #[allow(unused)]
153        fn func(a: i32, b: i32) {}
154        Function::new_native(store, func)
155    }
156
157    pub fn log_component(store: &Store) -> Function {
158        #[allow(unused)]
159        fn func(a: i32) {}
160        Function::new_native(store, func)
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167    use std::{collections::HashMap, path::PathBuf};
168
169    struct TestCase<'a> {
170        circuit_path: &'a str,
171        inputs_path: &'a str,
172        n_vars: u32,
173        n64: u32,
174        witness: &'a [&'a str],
175    }
176
177    pub fn root_path(p: &str) -> String {
178        let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
179        path.push(p);
180        path.to_string_lossy().to_string()
181    }
182
183    #[test]
184    fn multiplier_1() {
185        run_test(TestCase {
186            circuit_path: root_path("test-vectors/mycircuit.wasm").as_str(),
187            inputs_path: root_path("test-vectors/mycircuit-input1.json").as_str(),
188            n_vars: 4,
189            n64: 4,
190            witness: &["1", "33", "3", "11"],
191        });
192    }
193
194    #[test]
195    fn multiplier_2() {
196        run_test(TestCase {
197            circuit_path: root_path("test-vectors/mycircuit.wasm").as_str(),
198            inputs_path: root_path("test-vectors/mycircuit-input2.json").as_str(),
199            n_vars: 4,
200            n64: 4,
201            witness: &[
202                "1",
203                "21888242871839275222246405745257275088548364400416034343698204186575672693159",
204                "21888242871839275222246405745257275088548364400416034343698204186575796149939",
205                "11",
206            ],
207        });
208    }
209
210    #[test]
211    fn multiplier_3() {
212        run_test(TestCase {
213            circuit_path: root_path("test-vectors/mycircuit.wasm").as_str(),
214            inputs_path: root_path("test-vectors/mycircuit-input3.json").as_str(),
215            n_vars: 4,
216            n64: 4,
217            witness: &[
218                "1",
219                "21888242871839275222246405745257275088548364400416034343698204186575808493616",
220                "10944121435919637611123202872628637544274182200208017171849102093287904246808",
221                "2",
222            ],
223        });
224    }
225
226    #[test]
227    fn safe_multipler() {
228        let witness =
229            std::fs::read_to_string(&root_path("test-vectors/safe-circuit-witness.json")).unwrap();
230        let witness: Vec<String> = serde_json::from_str(&witness).unwrap();
231        let witness = &witness.iter().map(|x| x.as_ref()).collect::<Vec<_>>();
232        run_test(TestCase {
233            circuit_path: root_path("test-vectors/circuit2.wasm").as_str(),
234            inputs_path: root_path("test-vectors/mycircuit-input1.json").as_str(),
235            n_vars: 132, // 128 + 4
236            n64: 4,
237            witness,
238        });
239    }
240
241    #[test]
242    fn smt_verifier() {
243        let witness =
244            std::fs::read_to_string(&root_path("test-vectors/smtverifier10-witness.json")).unwrap();
245        let witness: Vec<String> = serde_json::from_str(&witness).unwrap();
246        let witness = &witness.iter().map(|x| x.as_ref()).collect::<Vec<_>>();
247
248        run_test(TestCase {
249            circuit_path: root_path("test-vectors/smtverifier10.wasm").as_str(),
250            inputs_path: root_path("test-vectors/smtverifier10-input.json").as_str(),
251            n_vars: 4794,
252            n64: 4,
253            witness,
254        });
255    }
256
257    use serde_json::Value;
258    use std::str::FromStr;
259
260    fn value_to_bigint(v: Value) -> BigInt {
261        match v {
262            Value::String(inner) => BigInt::from_str(&inner).unwrap(),
263            Value::Number(inner) => BigInt::from(inner.as_u64().expect("not a u32")),
264            _ => panic!("unsupported type"),
265        }
266    }
267
268    fn run_test(case: TestCase) {
269        let mut wtns = WitnessCalculator::new(case.circuit_path).unwrap();
270        assert_eq!(
271            wtns.memory.prime.to_str_radix(16),
272            "30644E72E131A029B85045B68181585D2833E84879B9709143E1F593F0000001".to_lowercase()
273        );
274        assert_eq!(wtns.instance.get_n_vars().unwrap() as u32, case.n_vars);
275        assert_eq!(wtns.n64 as u32, case.n64);
276
277        let inputs_str = std::fs::read_to_string(case.inputs_path).unwrap();
278        let inputs: std::collections::HashMap<String, serde_json::Value> =
279            serde_json::from_str(&inputs_str).unwrap();
280
281        let inputs = inputs
282            .iter()
283            .map(|(key, value)| {
284                let res = match value {
285                    Value::String(inner) => {
286                        vec![BigInt::from_str(inner).unwrap()]
287                    }
288                    Value::Number(inner) => {
289                        vec![BigInt::from(inner.as_u64().expect("not a u32"))]
290                    }
291                    Value::Array(inner) => inner.iter().cloned().map(value_to_bigint).collect(),
292                    _ => panic!(),
293                };
294
295                (key.clone(), res)
296            })
297            .collect::<HashMap<_, _>>();
298
299        let res = wtns.calculate_witness(inputs, false).unwrap();
300        for (r, w) in res.iter().zip(case.witness) {
301            assert_eq!(r, &BigInt::from_str(w).unwrap());
302        }
303    }
304}