1use 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
30pub 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
52pub 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
62pub 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 if expect_satisfied {
84 res.unwrap();
85 } else {
86 assert!(res.is_err());
87 }
88}
89
90pub 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
101pub 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
118pub fn base_test() -> BaseTester {
120 BaseTester::default()
121}
122
123impl BaseTester {
124 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 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 pub fn expect_satisfied(mut self, expect_satisfied: bool) -> Self {
141 self.expect_satisfied = expect_satisfied;
142 self
143 }
144
145 pub fn unusable_rows(mut self, unusable_rows: usize) -> Self {
147 self.unusable_rows = unusable_rows;
148 self
149 }
150
151 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 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 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 let res = f(builder.pool(0), &range);
175
176 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 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 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 f(builder.pool(0), &range, init_input);
212
213 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 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(¶ms, &builder).unwrap();
225 end_timer!(vk_time);
226 let pk_time = start_timer!(|| "Generating pkey");
227 let pk = keygen_pk(¶ms, vk, &builder).unwrap();
228 end_timer!(pk_time);
229
230 let break_points = builder.break_points();
231 drop(builder);
232 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(¶ms, &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(¶ms, 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
250pub struct BenchStats {
252 pub config_params: BaseCircuitParams,
254 pub vk_time: TimerInfo,
256 pub pk_time: TimerInfo,
258 pub proof_time: TimerInfo,
260 pub proof_size: usize,
262 pub verify_time: TimerInfo,
264}