Skip to main content

halo2_base/utils/
testing.rs

1//! Utilities for testing
2use crate::{
3    gates::{
4        circuit::{builder::RangeCircuitBuilder, BaseCircuitParams, CircuitBuilderStage},
5        flex_gate::threads::SinglePhaseCoreManager,
6        GateChip, RangeChip,
7    },
8    halo2_proofs::{
9        dev::MockProver,
10        halo2curves::bn256::{Bn256, Fr, G1Affine},
11        plonk::{
12            create_proof, keygen_pk, keygen_vk, verify_proof, Circuit, ProvingKey, VerifyingKey,
13        },
14        poly::commitment::ParamsProver,
15        poly::kzg::{
16            commitment::KZGCommitmentScheme, commitment::ParamsKZG, multiopen::ProverSHPLONK,
17            multiopen::VerifierSHPLONK, strategy::SingleStrategy,
18        },
19        transcript::{
20            Blake2bRead, Blake2bWrite, Challenge255, TranscriptReadBuffer, TranscriptWriterBuffer,
21        },
22    },
23    Context,
24};
25use ark_std::{end_timer, perf_trace::TimerInfo, start_timer};
26use rand::{rngs::StdRng, SeedableRng};
27
28use super::fs::gen_srs;
29
30/// Helper function to generate a proof with real prover using SHPLONK KZG multi-open polynomical commitment scheme
31/// and Blake2b as the hash function for Fiat-Shamir.
32pub fn gen_proof_with_instances(
33    params: &ParamsKZG<Bn256>,
34    pk: &ProvingKey<G1Affine>,
35    circuit: impl Circuit<Fr>,
36    instances: &[&[Fr]],
37) -> Vec<u8> {
38    let rng = StdRng::seed_from_u64(0);
39    let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
40    create_proof::<
41        KZGCommitmentScheme<Bn256>,
42        ProverSHPLONK<'_, Bn256>,
43        Challenge255<_>,
44        _,
45        Blake2bWrite<Vec<u8>, G1Affine, _>,
46        _,
47    >(params, pk, &[circuit], &[instances], rng, &mut transcript)
48    .expect("prover should not fail");
49    transcript.finalize()
50}
51
52/// For testing use only: Helper function to generate a proof **without public instances** with real prover using SHPLONK KZG multi-open polynomical commitment scheme
53/// and Blake2b as the hash function for Fiat-Shamir.
54pub fn gen_proof(
55    params: &ParamsKZG<Bn256>,
56    pk: &ProvingKey<G1Affine>,
57    circuit: impl Circuit<Fr>,
58) -> Vec<u8> {
59    gen_proof_with_instances(params, pk, circuit, &[])
60}
61
62/// Helper function to verify a proof (generated using [`gen_proof_with_instances`]) using SHPLONK KZG multi-open polynomical commitment scheme
63/// and Blake2b as the hash function for Fiat-Shamir.
64pub fn check_proof_with_instances(
65    params: &ParamsKZG<Bn256>,
66    vk: &VerifyingKey<G1Affine>,
67    proof: &[u8],
68    instances: &[&[Fr]],
69    expect_satisfied: bool,
70) {
71    let verifier_params = params.verifier_params();
72    let strategy = SingleStrategy::new(params);
73    let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(proof);
74    let res = verify_proof::<
75        KZGCommitmentScheme<Bn256>,
76        VerifierSHPLONK<'_, Bn256>,
77        Challenge255<G1Affine>,
78        Blake2bRead<&[u8], G1Affine, Challenge255<G1Affine>>,
79        SingleStrategy<'_, Bn256>,
80    >(verifier_params, vk, strategy, &[instances], &mut transcript);
81    // Just FYI, because strategy is `SingleStrategy`, the output `res` is `Result<(), Error>`, so there is no need to call `res.finalize()`.
82
83    if expect_satisfied {
84        res.unwrap();
85    } else {
86        assert!(res.is_err());
87    }
88}
89
90/// For testing only: Helper function to verify a proof (generated using [`gen_proof`]) without public instances using SHPLONK KZG multi-open polynomical commitment scheme
91/// and Blake2b as the hash function for Fiat-Shamir.
92pub fn check_proof(
93    params: &ParamsKZG<Bn256>,
94    vk: &VerifyingKey<G1Affine>,
95    proof: &[u8],
96    expect_satisfied: bool,
97) {
98    check_proof_with_instances(params, vk, proof, &[], expect_satisfied);
99}
100
101/// Helper to facilitate easier writing of tests using `RangeChip` and `RangeCircuitBuilder`.
102/// By default, the [`MockProver`] is used.
103///
104/// Currently this tester uses all private inputs.
105pub struct BaseTester {
106    k: u32,
107    lookup_bits: Option<usize>,
108    expect_satisfied: bool,
109    unusable_rows: usize,
110}
111
112impl Default for BaseTester {
113    fn default() -> Self {
114        Self { k: 10, lookup_bits: Some(9), expect_satisfied: true, unusable_rows: 9 }
115    }
116}
117
118/// Creates a [`BaseTester`]
119pub fn base_test() -> BaseTester {
120    BaseTester::default()
121}
122
123impl BaseTester {
124    /// Changes the number of rows in the circuit to 2<sup>k</sup>.
125    /// By default it will also set lookup bits as large as possible, to `k - 1`.
126    pub fn k(mut self, k: u32) -> Self {
127        self.k = k;
128        self.lookup_bits = Some(k as usize - 1);
129        self
130    }
131
132    /// Sets the size of the lookup table used for range checks to [0, 2<sup>lookup_bits</sup>)
133    pub fn lookup_bits(mut self, lookup_bits: usize) -> Self {
134        assert!(lookup_bits < self.k as usize, "lookup_bits must be less than k");
135        self.lookup_bits = Some(lookup_bits);
136        self
137    }
138
139    /// Specify whether you expect this test to pass or fail. Default: pass
140    pub fn expect_satisfied(mut self, expect_satisfied: bool) -> Self {
141        self.expect_satisfied = expect_satisfied;
142        self
143    }
144
145    /// Set the number of blinding (poisoned) rows
146    pub fn unusable_rows(mut self, unusable_rows: usize) -> Self {
147        self.unusable_rows = unusable_rows;
148        self
149    }
150
151    /// Run a mock test by providing a closure that uses a `ctx` and `RangeChip`.
152    /// - `expect_satisfied`: flag for whether you expect the test to pass or fail. Failure means a constraint system failure -- the tester does not catch system panics.
153    pub fn run<R>(&self, f: impl FnOnce(&mut Context<Fr>, &RangeChip<Fr>) -> R) -> R {
154        self.run_builder(|builder, range| f(builder.main(), range))
155    }
156
157    /// Run a mock test by providing a closure that uses a `ctx` and `GateChip`.
158    /// - `expect_satisfied`: flag for whether you expect the test to pass or fail. Failure means a constraint system failure -- the tester does not catch system panics.
159    pub fn run_gate<R>(&self, f: impl FnOnce(&mut Context<Fr>, &GateChip<Fr>) -> R) -> R {
160        self.run(|ctx, range| f(ctx, &range.gate))
161    }
162
163    /// Run a mock test by providing a closure that uses a `builder` and `RangeChip`.
164    pub fn run_builder<R>(
165        &self,
166        f: impl FnOnce(&mut SinglePhaseCoreManager<Fr>, &RangeChip<Fr>) -> R,
167    ) -> R {
168        let mut builder = RangeCircuitBuilder::default().use_k(self.k as usize);
169        if let Some(lb) = self.lookup_bits {
170            builder.set_lookup_bits(lb)
171        }
172        let range = RangeChip::new(self.lookup_bits.unwrap_or(0), builder.lookup_manager().clone());
173        // run the function, mutating `builder`
174        let res = f(builder.pool(0), &range);
175
176        // helper check: if your function didn't use lookups, turn lookup table "off"
177        let t_cells_lookup =
178            builder.lookup_manager().iter().map(|lm| lm.total_rows()).sum::<usize>();
179        let lookup_bits = if t_cells_lookup == 0 { None } else { self.lookup_bits };
180        builder.config_params.lookup_bits = lookup_bits;
181
182        // configure the circuit shape, 9 blinding rows seems enough
183        builder.calculate_params(Some(self.unusable_rows));
184        if self.expect_satisfied {
185            MockProver::run(self.k, &builder, vec![]).unwrap().assert_satisfied();
186        } else {
187            assert!(MockProver::run(self.k, &builder, vec![]).unwrap().verify().is_err());
188        }
189        res
190    }
191
192    /// Runs keygen, real prover, and verifier by providing a closure that uses a `builder` and `RangeChip`.
193    ///
194    /// Must provide `init_input` for use during key generation, which is preferably not equal to `logic_input`.
195    /// These are the inputs to the closure, not necessary public inputs to the circuit.
196    ///
197    /// Currently for testing, no public instances.
198    pub fn bench_builder<I: Clone>(
199        &self,
200        init_input: I,
201        logic_input: I,
202        f: impl Fn(&mut SinglePhaseCoreManager<Fr>, &RangeChip<Fr>, I),
203    ) -> BenchStats {
204        let mut builder =
205            RangeCircuitBuilder::from_stage(CircuitBuilderStage::Keygen).use_k(self.k as usize);
206        if let Some(lb) = self.lookup_bits {
207            builder.set_lookup_bits(lb)
208        }
209        let range = RangeChip::new(self.lookup_bits.unwrap_or(0), builder.lookup_manager().clone());
210        // run the function, mutating `builder`
211        f(builder.pool(0), &range, init_input);
212
213        // helper check: if your function didn't use lookups, turn lookup table "off"
214        let t_cells_lookup =
215            builder.lookup_manager().iter().map(|lm| lm.total_rows()).sum::<usize>();
216        let lookup_bits = if t_cells_lookup == 0 { None } else { self.lookup_bits };
217        builder.config_params.lookup_bits = lookup_bits;
218
219        // configure the circuit shape, 9 blinding rows seems enough
220        let config_params = builder.calculate_params(Some(self.unusable_rows));
221
222        let params = gen_srs(self.k);
223        let vk_time = start_timer!(|| "Generating vkey");
224        let vk = keygen_vk(&params, &builder).unwrap();
225        end_timer!(vk_time);
226        let pk_time = start_timer!(|| "Generating pkey");
227        let pk = keygen_pk(&params, vk, &builder).unwrap();
228        end_timer!(pk_time);
229
230        let break_points = builder.break_points();
231        drop(builder);
232        // create real proof
233        let proof_time = start_timer!(|| "Proving time");
234        let mut builder = RangeCircuitBuilder::prover(config_params.clone(), break_points);
235        let range = RangeChip::new(self.lookup_bits.unwrap_or(0), builder.lookup_manager().clone());
236        f(builder.pool(0), &range, logic_input);
237        let proof = gen_proof(&params, &pk, builder);
238        end_timer!(proof_time);
239
240        let proof_size = proof.len();
241
242        let verify_time = start_timer!(|| "Verify time");
243        check_proof(&params, pk.get_vk(), &proof, self.expect_satisfied);
244        end_timer!(verify_time);
245
246        BenchStats { config_params, vk_time, pk_time, proof_time, proof_size, verify_time }
247    }
248}
249
250/// Bench stats
251pub struct BenchStats {
252    /// Config params
253    pub config_params: BaseCircuitParams,
254    /// Vkey gen time
255    pub vk_time: TimerInfo,
256    /// Pkey gen time
257    pub pk_time: TimerInfo,
258    /// Proving time
259    pub proof_time: TimerInfo,
260    /// Proof size in bytes
261    pub proof_size: usize,
262    /// Verify time
263    pub verify_time: TimerInfo,
264}