circuits/test/utils/
circom_tester.rs

1use 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
16/// Represents a Circom-style hierarchical signal path.
17impl SignalKey {
18    /// Creates a new base signal key.
19    pub fn new(base: impl Into<String>) -> Self {
20        Self(base.into())
21    }
22
23    /// Appends an array index (`[...]`) to the key.
24    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    /// Appends a field accessor (`.field`) to the key.
32    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
45/// Allow common types to be converted into InputValue.
46impl 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/// Storage for Circom input signals.
83/// Wraps a hashmap of `String → InputValue`.
84///
85/// Example:
86///
87/// ```
88/// use circuits::test::utils::circom_tester::{Inputs, SignalKey};
89/// use zkhash::fields::bn256::FpBN256 as Scalar;
90/// let mut inputs = Inputs::new();
91/// inputs.set("root", Scalar::from(5));
92/// inputs.set_key(&SignalKey::new("arr").idx(0), Scalar::from(10));
93/// ```
94#[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    /// Sets an input using a plain string key.
107    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    /// Set using a SignalKey path (e.g., membershipProofs[0][0].leaf).
116    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/// Represents a single Circom input value
129#[derive(Clone, Debug)]
130pub enum InputValue {
131    Single(BigInt),
132    Array(Vec<BigInt>),
133}
134
135/// Contains the Groth16 proving key (pk),
136/// verifying key (vk), and the *processed* verifying key (pvk).
137#[derive(Clone)]
138pub struct CircuitKeys {
139    pub pk: ProvingKey<Bn254>,
140    pub vk: VerifyingKey<Bn254>,
141    pub pvk: PreparedVerifyingKey<Bn254>,
142}
143
144/// Result of proving + verifying a Circom circuit
145#[derive(Clone, Debug)]
146pub struct CircomResult {
147    pub verified: bool,
148    pub public_inputs: Vec<Fr>, /* this can be a trait but we dont care about generalising that
149                                 * much now */
150    pub proof: Proof<Bn254>,
151    pub vk: VerifyingKey<Bn254>,
152}
153
154/// Generates Groth16 proving + verifying keys for a Circom circuit.
155/// This operation is expensive and should be done once when testing
156/// many input combinations.
157pub 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    // No inputs: just the empty circuit for setup
167    let empty = builder.setup();
168    let mut rng = thread_rng();
169
170    // Use default LibsnarkReduction for WASM prover compatibility
171    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
179/// Loads Groth16 keys from a binary proving key file.
180///
181/// The proving key file should be serialized using
182/// `ark_serialize::CanonicalSerialize`. The verification key is extracted from
183/// the proving key, and the prepared verification key is computed for efficient
184/// verification.
185///
186/// # Arguments
187///
188/// * `pk_path` - Path to the binary proving key file
189///
190/// # Returns
191///
192/// Returns `Ok(CircuitKeys)` containing the proving key, verification key,
193/// and prepared verification key, or an error if loading fails.
194pub 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    // Extract verification key from proving key
203    let vk = pk.vk.clone();
204
205    // Compute prepared verification key for efficient verification
206    // Use default LibsnarkReduction for WASM prover compatibility
207    let pvk = Groth16::<Bn254>::process_vk(&vk).map_err(|e| anyhow!("process_vk failed: {e}"))?;
208
209    Ok(CircuitKeys { pk, vk, pvk })
210}
211
212/// Proves and verifies a Circom circuit using precomputed Groth16 keys.
213/// This is the preferred function when repeated proofs must be generated.
214///
215/// Steps:
216/// 1. Load Circom config (WASM + R1CS)
217/// 2. Build circuit with provided inputs
218/// 3. Generate Groth16 proof using precomputed `pk`
219/// 4. Verify the proof using fast `pvk`
220///
221/// Returns `CircomResult`.
222pub 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    // Use default LibsnarkReduction for WASM prover compatibility
242    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
260/// Internal helper for adding input values into the Circom builder.
261/// Arrays are pushed element-by-element.
262fn 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
275/// Proves and verifies a Circom circuit, generating keys on each call
276///
277/// Convenience function that generates Groth16 keys and then proves and
278/// verifies the circuit. This is simpler to use but less efficient for repeated
279/// proofs since key generation is expensive. For multiple proofs with the same
280/// circuit, use `generate_keys` once and then call `prove_and_verify_with_keys`
281/// repeatedly.
282///
283/// # Arguments
284///
285/// * `wasm_path` - Path to the compiled WASM file for witness generation
286/// * `r1cs_path` - Path to the R1CS constraint system file
287/// * `inputs` - Circuit input values to use for proving
288///
289/// # Returns
290///
291/// Returns `Ok(CircomResult)` containing the verification result, proof, public
292/// inputs, and verifying key, or an error if key generation, proving, or
293/// verification fails.
294pub 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}