circuits/test/utils/
transaction_case.rs

1use super::{
2    circom_tester::prove_and_verify,
3    general::scalar_to_bigint,
4    keypair::{derive_public_key, sign},
5    merkle_tree::{merkle_proof, merkle_root},
6    transaction::{commitment, nullifier},
7};
8use crate::test::utils::circom_tester::Inputs;
9use anyhow::{Result, ensure};
10use num_bigint::BigInt;
11use std::{
12    panic::{self, AssertUnwindSafe},
13    path::PathBuf,
14};
15use zkhash::fields::bn256::FpBN256 as Scalar;
16
17#[derive(Clone, Debug)]
18/// Description of a note spent by the tested transaction.
19pub struct InputNote {
20    pub leaf_index: usize, /* We need to place the note in the tree, and hold the index to know
21                            * where it is */
22    pub priv_key: Scalar, // Used to derive its public key and to sign nullifiers for spends.
23    pub blinding: Scalar, // Keeps the commitment hiding so tests match the production circuit.
24    pub amount: Scalar,   // Amount being spent; required for balance and commitment inputs.
25}
26
27#[derive(Clone, Debug)]
28/// Description of a note created by the tested transaction.
29pub struct OutputNote {
30    pub pub_key: Scalar,
31    pub blinding: Scalar,
32    pub amount: Scalar,
33}
34
35#[derive(Clone, Debug)]
36/// Convenience container holding a single test transaction scenario.
37/// We use `Vec` because we usually have more than one input and output. The
38/// test defines how many
39pub struct TxCase {
40    pub inputs: Vec<InputNote>,
41    pub outputs: Vec<OutputNote>,
42}
43
44impl TxCase {
45    pub fn new(inputs: Vec<InputNote>, outputs: Vec<OutputNote>) -> Self {
46        Self { inputs, outputs }
47    }
48}
49
50pub struct TransactionWitness {
51    pub root: Scalar,
52    pub public_keys: Vec<Scalar>,
53    pub nullifiers: Vec<Scalar>,
54    pub path_indices: Vec<Scalar>,
55    pub path_elements_flat: Vec<BigInt>,
56}
57
58/// Builds the witnesses needed to exercise a `TxCase`
59///
60/// Populates commitment leaves in the Merkle tree, derives Merkle proofs for
61/// each input note, and computes nullifiers. This prepares all the witness data
62/// required for proving a transaction.
63///
64/// # Arguments
65///
66/// * `case` - Transaction case containing input and output notes
67/// * `leaves` - Initial leaves vector (will be modified with commitments)
68/// * `expected_levels` - Expected number of levels in the Merkle tree
69///
70/// # Returns
71///
72/// Returns `Ok(TransactionWitness)` containing the root, public keys,
73/// nullifiers, path indices, and flattened path elements, or an error if the
74/// tree depth doesn't match expectations.
75pub fn prepare_transaction_witness(
76    case: &TxCase,
77    mut leaves: Vec<Scalar>,
78    expected_levels: usize,
79) -> Result<TransactionWitness> {
80    let mut commitments = Vec::with_capacity(case.inputs.len());
81    let mut public_keys = Vec::with_capacity(case.inputs.len());
82
83    for note in &case.inputs {
84        let pk = derive_public_key(note.priv_key);
85        let cm = commitment(note.amount, pk, note.blinding);
86        public_keys.push(pk);
87        commitments.push(cm);
88        leaves[note.leaf_index] = cm;
89    }
90
91    let root = merkle_root(leaves.clone());
92    let mut path_indices = Vec::with_capacity(case.inputs.len());
93    let mut path_elements_flat =
94        Vec::with_capacity(expected_levels.saturating_mul(case.inputs.len()));
95    let mut nullifiers = Vec::with_capacity(case.inputs.len());
96
97    for (i, note) in case.inputs.iter().enumerate() {
98        let (siblings, path_idx_u64, depth) = merkle_proof(&leaves, note.leaf_index);
99        ensure!(
100            depth == expected_levels,
101            "unexpected depth for input {i}, expected {expected_levels}, got {depth}"
102        );
103
104        // Flatten sibling nodes into the format the Circom tester expects.
105        path_elements_flat.extend(siblings.into_iter().map(scalar_to_bigint));
106
107        let path_idx = Scalar::from(path_idx_u64);
108        path_indices.push(path_idx);
109
110        let sig = sign(note.priv_key, commitments[i], path_idx);
111        let nul = nullifier(commitments[i], path_idx, sig);
112        nullifiers.push(nul);
113    }
114
115    Ok(TransactionWitness {
116        root,
117        public_keys,
118        nullifiers,
119        path_indices,
120        path_elements_flat,
121    })
122}
123
124/// Populates Circom tester inputs for policy-enabled and regular transactions
125///
126/// Builds the input structure required by the Circom circuit tester from a
127/// transaction case and its witness data. Includes all public and private
128/// inputs needed for proving.
129///
130/// # Arguments
131///
132/// * `case` - Transaction case containing input and output notes
133/// * `witness` - Transaction witness containing Merkle proofs and nullifiers
134/// * `public_amount` - Public amount scalar value (net public input/output)
135///
136/// # Returns
137///
138/// Returns an `Inputs` structure populated with all circuit inputs.
139pub fn build_base_inputs(
140    case: &TxCase,
141    witness: &TransactionWitness,
142    public_amount: Scalar,
143) -> Inputs {
144    let mut inputs = Inputs::new();
145
146    inputs.set("root", scalar_to_bigint(witness.root));
147    inputs.set("publicAmount", scalar_to_bigint(public_amount));
148    inputs.set("extDataHash", BigInt::from(0u32));
149
150    inputs.set("inputNullifier", witness.nullifiers.clone());
151    inputs.set(
152        "inAmount",
153        case.inputs
154            .iter()
155            .map(|n| n.amount)
156            .collect::<Vec<Scalar>>(),
157    );
158    inputs.set(
159        "inPrivateKey",
160        case.inputs
161            .iter()
162            .map(|n| n.priv_key)
163            .collect::<Vec<Scalar>>(),
164    );
165    inputs.set(
166        "inBlinding",
167        case.inputs
168            .iter()
169            .map(|n| n.blinding)
170            .collect::<Vec<Scalar>>(),
171    );
172    inputs.set("inPathIndices", witness.path_indices.clone());
173    inputs.set("inPathElements", witness.path_elements_flat.clone());
174
175    let output_commitments: Vec<BigInt> = case
176        .outputs
177        .iter()
178        .map(|out| scalar_to_bigint(commitment(out.amount, out.pub_key, out.blinding)))
179        .collect();
180    inputs.set("outputCommitment", output_commitments);
181
182    inputs.set(
183        "outAmount",
184        case.outputs
185            .iter()
186            .map(|n| n.amount)
187            .collect::<Vec<Scalar>>(),
188    );
189    inputs.set(
190        "outPubkey",
191        case.outputs
192            .iter()
193            .map(|n| n.pub_key)
194            .collect::<Vec<Scalar>>(),
195    );
196    inputs.set(
197        "outBlinding",
198        case.outputs
199            .iter()
200            .map(|n| n.blinding)
201            .collect::<Vec<Scalar>>(),
202    );
203
204    inputs
205}
206
207/// Runs a Circom proof/verify cycle for a transaction test case
208///
209/// Prepares the transaction witness, builds circuit inputs, and executes
210/// a proof generation and verification cycle.
211///
212/// # Arguments
213///
214/// * `wasm` - Path to the compiled WASM file
215/// * `r1cs` - Path to the R1CS constraint system file
216/// * `case` - Transaction case to prove
217/// * `leaves` - Initial leaves vector for the Merkle tree
218/// * `public_amount` - Public amount scalar value
219/// * `expected_levels` - Expected number of levels in the Merkle tree
220///
221/// # Returns
222///
223/// Returns `Ok(())` if the proof is generated and verified successfully,
224/// or an error if witness preparation, proving, or verification fails.
225pub fn prove_transaction_case(
226    wasm: &PathBuf,
227    r1cs: &PathBuf,
228    case: &TxCase,
229    leaves: Vec<Scalar>,
230    public_amount: Scalar,
231    expected_levels: usize,
232) -> Result<()> {
233    let witness = prepare_transaction_witness(case, leaves, expected_levels)?;
234    let inputs = build_base_inputs(case, &witness, public_amount);
235
236    let prove_result =
237        panic::catch_unwind(AssertUnwindSafe(|| prove_and_verify(wasm, r1cs, &inputs)));
238
239    match prove_result {
240        Ok(Ok(res)) if res.verified => Ok(()),
241        Ok(Ok(_)) => Err(anyhow::anyhow!(
242            "Proof failed to verify (res.verified=false)"
243        )),
244        Ok(Err(e)) => Err(anyhow::anyhow!("Prover error: {e:?}")),
245        Err(panic_info) => {
246            // Tests expect panics for invalid proofs; convert any panic into a typed error.
247            let msg = if let Some(s) = panic_info.downcast_ref::<&str>() {
248                s.to_string()
249            } else if let Some(s) = panic_info.downcast_ref::<String>() {
250                s.clone()
251            } else {
252                "Unknown panic".to_string()
253            };
254            Err(anyhow::anyhow!(
255                "Prover panicked (expected on invalid proof): {msg}"
256            ))
257        }
258    }
259}