circuits/test/utils/
circom_tester.rs1use super::general::scalar_to_bigint;
2use anyhow::{Context, Result, anyhow};
3use ark_bn254::{Bn254, Fr};
4use ark_circom::{CircomBuilder, CircomConfig};
5use ark_groth16::{Groth16, PreparedVerifyingKey, Proof, ProvingKey, VerifyingKey};
6use ark_serialize::CanonicalDeserialize;
7use ark_snark::SNARK;
8use ark_std::rand::thread_rng;
9use num_bigint::BigInt;
10use std::{collections::HashMap, fmt, fmt::Display, fs::File, io::BufReader, path::Path};
11use zkhash::fields::bn256::FpBN256 as Scalar;
12
13#[derive(Clone, Debug)]
14pub struct SignalKey(String);
15
16impl SignalKey {
18 pub fn new(base: impl Into<String>) -> Self {
20 Self(base.into())
21 }
22
23 pub fn idx(mut self, i: usize) -> Self {
25 self.0.push('[');
26 self.0.push_str(&i.to_string());
27 self.0.push(']');
28 self
29 }
30
31 pub fn field(mut self, name: &str) -> Self {
33 self.0.push('.');
34 self.0.push_str(name);
35 self
36 }
37}
38
39impl Display for SignalKey {
40 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
41 f.write_str(&self.0)
42 }
43}
44
45impl From<BigInt> for InputValue {
47 fn from(value: BigInt) -> Self {
48 InputValue::Single(value)
49 }
50}
51
52impl From<&BigInt> for InputValue {
53 fn from(value: &BigInt) -> Self {
54 InputValue::Single(value.clone())
55 }
56}
57
58impl From<Vec<BigInt>> for InputValue {
59 fn from(value: Vec<BigInt>) -> Self {
60 InputValue::Array(value)
61 }
62}
63
64impl From<Scalar> for InputValue {
65 fn from(value: Scalar) -> Self {
66 InputValue::Single(scalar_to_bigint(value))
67 }
68}
69
70impl From<&Scalar> for InputValue {
71 fn from(value: &Scalar) -> Self {
72 InputValue::Single(scalar_to_bigint(*value))
73 }
74}
75
76impl From<Vec<Scalar>> for InputValue {
77 fn from(values: Vec<Scalar>) -> Self {
78 InputValue::Array(values.into_iter().map(scalar_to_bigint).collect())
79 }
80}
81
82#[derive(Default)]
95pub struct Inputs {
96 inner: HashMap<String, InputValue>,
97}
98
99impl Inputs {
100 pub fn new() -> Self {
101 Self {
102 inner: HashMap::new(),
103 }
104 }
105
106 pub fn set<K, V>(&mut self, key: K, value: V)
108 where
109 K: Into<String>,
110 V: Into<InputValue>,
111 {
112 self.inner.insert(key.into(), value.into());
113 }
114
115 pub fn set_key<V>(&mut self, key: &SignalKey, value: V)
117 where
118 V: Into<InputValue>,
119 {
120 self.inner.insert(key.to_string(), value.into());
121 }
122
123 pub fn iter(&self) -> impl Iterator<Item = (&String, &InputValue)> {
124 self.inner.iter()
125 }
126}
127
128#[derive(Clone, Debug)]
130pub enum InputValue {
131 Single(BigInt),
132 Array(Vec<BigInt>),
133}
134
135#[derive(Clone)]
138pub struct CircuitKeys {
139 pub pk: ProvingKey<Bn254>,
140 pub vk: VerifyingKey<Bn254>,
141 pub pvk: PreparedVerifyingKey<Bn254>,
142}
143
144#[derive(Clone, Debug)]
146pub struct CircomResult {
147 pub verified: bool,
148 pub public_inputs: Vec<Fr>, pub proof: Proof<Bn254>,
151 pub vk: VerifyingKey<Bn254>,
152}
153
154pub fn generate_keys(
158 wasm_path: impl AsRef<Path>,
159 r1cs_path: impl AsRef<Path>,
160) -> Result<CircuitKeys> {
161 let cfg = CircomConfig::<Fr>::new(wasm_path.as_ref(), r1cs_path.as_ref())
162 .map_err(|e| anyhow!("CircomConfig error: {e}"))?;
163
164 let builder = CircomBuilder::new(cfg);
165
166 let empty = builder.setup();
168 let mut rng = thread_rng();
169
170 let (pk, vk) = Groth16::<Bn254>::circuit_specific_setup(empty, &mut rng)
172 .map_err(|e| anyhow!("circuit_specific_setup failed: {e}"))?;
173
174 let pvk = Groth16::<Bn254>::process_vk(&vk).map_err(|e| anyhow!("process_vk failed: {e}"))?;
175
176 Ok(CircuitKeys { pk, vk, pvk })
177}
178
179pub fn load_keys(pk_path: impl AsRef<Path>) -> Result<CircuitKeys> {
195 let file = File::open(pk_path.as_ref())
196 .with_context(|| format!("Failed to open proving key file: {:?}", pk_path.as_ref()))?;
197 let mut reader = BufReader::new(file);
198
199 let pk: ProvingKey<Bn254> = ProvingKey::deserialize_compressed(&mut reader)
200 .map_err(|e| anyhow!("Failed to deserialize proving key: {e}"))?;
201
202 let vk = pk.vk.clone();
204
205 let pvk = Groth16::<Bn254>::process_vk(&vk).map_err(|e| anyhow!("process_vk failed: {e}"))?;
208
209 Ok(CircuitKeys { pk, vk, pvk })
210}
211
212pub fn prove_and_verify_with_keys(
223 wasm_path: impl AsRef<Path>,
224 r1cs_path: impl AsRef<Path>,
225 inputs: &Inputs,
226 keys: &CircuitKeys,
227) -> Result<CircomResult> {
228 let cfg = CircomConfig::<Fr>::new(wasm_path.as_ref(), r1cs_path.as_ref())
229 .map_err(|e| anyhow!("CircomConfig error: {e}"))?;
230
231 let mut builder = CircomBuilder::new(cfg);
232
233 for (signal, value) in inputs.iter() {
234 push_value(&mut builder, signal, value);
235 }
236
237 let circuit = builder.build().map_err(|e| anyhow!("build failed: {e}"))?;
238
239 let mut rng = thread_rng();
240
241 let proof = Groth16::<Bn254>::prove(&keys.pk, circuit.clone(), &mut rng)
243 .map_err(|e| anyhow!("prove failed: {e}"))?;
244
245 let public_inputs = circuit
246 .get_public_inputs()
247 .ok_or_else(|| anyhow!("get_public_inputs returned None"))?;
248
249 let verified = Groth16::<Bn254>::verify_with_processed_vk(&keys.pvk, &public_inputs, &proof)
250 .map_err(|e| anyhow!("verify_with_processed_vk failed: {e}"))?;
251
252 Ok(CircomResult {
253 verified,
254 public_inputs,
255 proof,
256 vk: keys.vk.clone(),
257 })
258}
259
260fn push_value(builder: &mut CircomBuilder<Fr>, path: &str, value: &InputValue) {
263 match value {
264 InputValue::Single(v) => {
265 builder.push_input(path, v.clone());
266 }
267 InputValue::Array(arr) => {
268 for v in arr.iter() {
269 builder.push_input(path, v.clone())
270 }
271 }
272 }
273}
274
275pub fn prove_and_verify(
295 wasm_path: impl AsRef<Path>,
296 r1cs_path: impl AsRef<Path>,
297 inputs: &Inputs,
298) -> Result<CircomResult> {
299 let keys = generate_keys(&wasm_path, &r1cs_path)?;
300 prove_and_verify_with_keys(wasm_path, r1cs_path, inputs, &keys)
301}