plonky2-recursion-experiment/src/tree_recursion.rs
src/tree_recursion.rsRaw
1/// N-arity tree of recursion with conditionals.
2///
3/// p_root
4/// ▲
5/// │
6/// ┌────────┐
7/// │ F │
8/// └────────┘
9/// ▲ ▲ ▲ ▲
10/// ┌─┘ │ │ └─┐
11/// ┌────┘ ┌─┘ └┐ └───┐
12/// │ │ ... │ │
13/// ┌────────┐┌┴┐┌─┐┌┴┐ ┌────────┐
14/// │ F ││.││.││.│ │ F │
15/// └────────┘└─┘└─┘└─┘ └────────┘
16/// ▲ ▲ ▲ ▲ ▲ ▲ ▲ ▲
17/// ┌─┘ │ └┐ └─┐ ┌─┘┌┘ └┐ └┐
18/// │ │ │ │ │ │ │ │
19/// p_1 p_2 ... p_n p'_1 p'_2... p'_n
20///
21///
22/// where each p_i is either
23/// - signature verification
24/// - recursive plonky2 proof (proof that verifies previous proof)
25/// (generated by `RecursiveCircuit::prove_step` method)
26/// in other words, each p_i is checking:
27/// `(signature proof OR recursive proof)`
28///
29/// Each node of the recursion tree, ie. each F, verifies the N incoming p_i's, that is
30/// `(signature proof OR recursive proof) AND ... AND (signature proof OR recursive proof)`
31/// and produces a new proof.
32///
33///
34/// For example, if N is set to N=2, then we work with a binary recursion tree:
35/// p_root
36/// ▲
37/// │
38/// ┌─┴─┐
39/// │ F │
40/// └───┘
41/// ▲ ▲
42/// ┌─┘ └─┐
43/// ┌───┘ └───┐
44/// │p_5 │p_6
45/// ┌─┴─┐ ┌─┴─┐
46/// │ F │ │ F │
47/// └───┘ └───┘
48/// ▲ ▲ ▲ ▲
49/// ┌─┘ └─┐ ┌─┘ └─┐
50/// │ │ │ │
51/// p_1 p_2 p_3 p_4
52///
53/// So that each node (F box) is verifying 2 p_i's, ie:
54/// `(signature proof OR recursive proof) AND (signature proof OR recursive proof)`
55///
56///
57/// With N=3, each node will be verifying 3 p_i's.
58/// `(signature proof OR recursive proof) AND (signature proof OR recursive proof) AND (signature proof OR recursive proof)`
59///
60///
61///
62/// Also, notice that if we set N=1, it is directly a linear chain of recursive proofs ('tree' of
63/// arity 1):
64/// ┌─┐ ┌─┐ ┌─┐ ┌─┐
65/// ─────►│F├────►│F├────►│F├────►│F├────►
66/// p_1 └─┘ p_2 └─┘ p_3 └─┘ p_4 └─┘ p_5
67///
68/// where each p_i is proving: `(signature proof OR recursive proof)`.
69///
70///
71/// To run the tests that checks this logic:
72/// cargo test --release test_tree_recursion -- --nocapture
73use anyhow::{anyhow, Result};
74use plonky2::field::types::Field;
75use plonky2::gates::noop::NoopGate;
76use plonky2::iop::witness::{PartialWitness, WitnessWrite};
77use plonky2::plonk::circuit_builder::CircuitBuilder;
78use plonky2::plonk::circuit_data::{
79 CircuitConfig, CircuitData, VerifierCircuitData, VerifierCircuitTarget,
80};
81use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
82use std::time::Instant;
83
84use sch::schnorr::*;
85use sch::schnorr_prover::*;
86
87use super::{sig_gadget::SignatureGadgetTargets, PlonkyProof, C, D, F};
88
89/// Contains the methods to `build` (ie. create the targets, the logic of the circuit), and
90/// `fill_targets` (ie. set the specific values to be used for the previously created targets).
91pub struct RecursiveCircuit<const N: usize> {
92 msg_targ: MessageTarget,
93 sigs_targ: Vec<SignatureGadgetTargets>,
94 proofs_targ: Vec<ProofWithPublicInputsTarget<D>>,
95 // the next two are common for all the given proofs. It is the data for this circuit itself
96 // (cyclic circuit).
97 verifier_data_targ: VerifierCircuitTarget,
98 verifier_data: VerifierCircuitData<F, C, D>,
99}
100
101impl<const N: usize> RecursiveCircuit<N> {
102 pub fn prepare_public_inputs(
103 verifier_data: VerifierCircuitData<F, C, D>,
104 msg: Vec<F>,
105 ) -> Vec<F> {
106 [
107 msg.clone(),
108 // add verifier_data as public inputs:
109 verifier_data.verifier_only.circuit_digest.elements.to_vec(),
110 verifier_data
111 .verifier_only
112 .constants_sigmas_cap
113 .0
114 .iter()
115 .flat_map(|e| e.elements)
116 .collect(),
117 ]
118 .concat()
119 }
120
121 // notice that this method does not fill the targets, which is done in the method
122 // `fill_recursive_circuit_targets`
123 pub fn build(
124 builder: &mut CircuitBuilder<F, D>,
125 verifier_data: VerifierCircuitData<F, C, D>,
126 msg_len: usize,
127 ) -> Result<Self> {
128 let msg_targ = MessageTarget::new_with_size(builder, msg_len);
129 // set msg as public input
130 builder.register_public_inputs(&msg_targ.msg);
131
132 // build the signature verification logic
133 let mut sigs_targ: Vec<SignatureGadgetTargets> = vec![];
134 for _ in 0..N {
135 let sig_targets = SignatureGadgetTargets::build(builder, &msg_targ)?;
136 sigs_targ.push(sig_targets);
137 }
138
139 // proof verification:
140
141 let common_data = verifier_data.common.clone();
142 let verifier_data_targ = builder.add_verifier_data_public_inputs();
143
144 let mut proofs_targ: Vec<ProofWithPublicInputsTarget<D>> = vec![];
145 for i in 0..N {
146 let proof_targ = builder.add_virtual_proof_with_pis(&common_data);
147 builder.conditionally_verify_cyclic_proof_or_dummy::<C>(
148 sigs_targ[i].selector_booltarg,
149 &proof_targ,
150 &common_data,
151 )?;
152 proofs_targ.push(proof_targ);
153 }
154
155 Ok(Self {
156 msg_targ,
157 sigs_targ,
158 proofs_targ,
159 verifier_data_targ,
160 verifier_data,
161 })
162 }
163
164 pub fn fill_targets(
165 &mut self,
166 pw: &mut PartialWitness<F>,
167 msg: &Vec<F>,
168 selectors: Vec<F>, // 1=proof, 0=sig
169 pks: &Vec<SchnorrPublicKey>,
170 sigs: &Vec<SchnorrSignature>,
171 recursive_proofs: &Vec<PlonkyProof>,
172 ) -> Result<()> {
173 // set the msg value (used by all N sig gadgets)
174 self.msg_targ.set_witness(pw, &msg).unwrap();
175
176 // set the signature related values
177 for i in 0..N {
178 self.sigs_targ[i].fill_targets(pw, selectors[i], &pks[i], &sigs[i])?;
179 }
180
181 // set proof related values:
182
183 // recursive proofs verification
184 pw.set_verifier_data_target(&self.verifier_data_targ, &self.verifier_data.verifier_only)?;
185
186 let public_inputs =
187 RecursiveCircuit::<N>::prepare_public_inputs(self.verifier_data.clone(), msg.clone());
188 for i in 0..N {
189 pw.set_proof_with_pis_target(
190 &self.proofs_targ[i],
191 &ProofWithPublicInputs {
192 proof: recursive_proofs[i].clone(),
193 public_inputs: public_inputs.clone(),
194 },
195 )?;
196 }
197
198 Ok(())
199 }
200}
201
202#[derive(Debug, Clone)]
203pub struct Recursion<const N: usize> {}
204
205pub fn common_data_for_recursion<const N: usize>(msg_len: usize) -> Result<CircuitData<F, C, D>> {
206 // 1st
207 let config = CircuitConfig::standard_recursion_config();
208 let builder = CircuitBuilder::<F, D>::new(config);
209 let data = builder.build::<C>();
210
211 // 2nd
212 let config = CircuitConfig::standard_recursion_config();
213 let mut builder = CircuitBuilder::<F, D>::new(config.clone());
214 let verifier_data = builder.add_virtual_verifier_data(data.common.config.fri_config.cap_height);
215 // proofs
216 for _ in 0..N {
217 let proof = builder.add_virtual_proof_with_pis(&data.common);
218 builder.verify_proof::<C>(&proof, &verifier_data, &data.common);
219 }
220 // let n_gates = builder.num_gates();
221 let data = builder.build::<C>();
222
223 // 3rd
224 let config = CircuitConfig::standard_recursion_config();
225 let mut builder = CircuitBuilder::<F, D>::new(config.clone());
226 let msg_targ = MessageTarget::new_with_size(&mut builder, msg_len);
227 // sigs verify
228 builder.register_public_inputs(&msg_targ.msg);
229
230 builder.add_gate(
231 // add a ConstantGate, because without this, when later generating the `dummy_circuit`
232 // (inside the `conditionally_verify_cyclic_proof_or_dummy`), it fails due the
233 // `CommonCircuitData` of the generated circuit not matching the given `CommonCircuitData`
234 // to create it. Without this it fails because it misses a ConstantGate.
235 plonky2::gates::constant::ConstantGate::new(config.num_constants),
236 vec![],
237 );
238
239 let _ = SignatureGadgetTargets::build(&mut builder, &msg_targ).unwrap();
240 let _ = SignatureGadgetTargets::build(&mut builder, &msg_targ).unwrap();
241
242 // proofs verify
243 let verifier_data = builder.add_verifier_data_public_inputs();
244 // proofs
245 for _ in 0..N {
246 let proof = builder.add_virtual_proof_with_pis(&data.common);
247 builder.verify_proof::<C>(&proof, &verifier_data, &data.common);
248 }
249
250 // pad min gates
251 let n_gates = compute_num_gates::<N>()?;
252 while builder.num_gates() < n_gates {
253 builder.add_gate(NoopGate, vec![]);
254 }
255 dbg!(builder.num_gates());
256 Ok(builder.build::<C>())
257}
258
259fn compute_num_gates<const N: usize>() -> Result<usize> {
260 // Note: the following numbers are WIP, obtained by trial-error by running different
261 // configurations in the tests.
262 let n_gates = match N {
263 1 => 1 << 12,
264 2 => 1 << 13,
265 3..=5 => 1 << 14,
266 6 => 1 << 15,
267 _ => 0,
268 };
269 if n_gates == 0 {
270 return Err(anyhow!(
271 "arity of N={} not supported yet. Currently supported N from 1 to 6 (both included)",
272 N
273 ));
274 }
275 Ok(n_gates)
276}
277
278impl<const N: usize> Recursion<N> {
279 /// returns the full-recursive CircuitData
280 pub fn circuit_data(msg_len: usize) -> Result<CircuitData<F, C, D>> {
281 let mut data = common_data_for_recursion::<N>(msg_len)?;
282
283 // build the actual RecursiveCircuit circuit data
284 let config = CircuitConfig::standard_recursion_config();
285 let mut builder = CircuitBuilder::new(config);
286 let _ = RecursiveCircuit::<N>::build(&mut builder, data.verifier_data(), msg_len)?;
287 dbg!(builder.num_gates());
288 data = builder.build::<C>();
289
290 Ok(data)
291 }
292
293 pub fn prove_step(
294 verifier_data: VerifierCircuitData<F, C, D>,
295 msg: &Vec<F>,
296 selectors: Vec<F>, // 1=proof, 0=sig
297 pks: &Vec<SchnorrPublicKey>,
298 sigs: &Vec<SchnorrSignature>,
299 recursive_proofs: &Vec<PlonkyProof>,
300 ) -> Result<PlonkyProof> {
301 println!("prove_step:");
302 for i in 0..N {
303 if selectors[i].is_nonzero() {
304 println!(" (selectors[{}]==1), verify {}-th proof", i, i);
305 } else {
306 println!(" (selectors[{}]==0), verify {}-th signature", i, i);
307 }
308 }
309
310 let config = CircuitConfig::standard_recursion_config();
311 let mut builder = CircuitBuilder::new(config);
312
313 // assign the targets
314 let start = Instant::now();
315 let mut circuit =
316 RecursiveCircuit::<N>::build(&mut builder, verifier_data.clone(), msg.len())?;
317 println!("RecursiveCircuit::build(): {:?}", start.elapsed());
318
319 // fill the targets
320 let mut pw = PartialWitness::new();
321 let start = Instant::now();
322 circuit.fill_targets(&mut pw, msg, selectors, pks, sigs, recursive_proofs)?;
323 println!("circuit.fill_targets(): {:?}", start.elapsed());
324
325 let start = Instant::now();
326 let data = builder.build::<C>();
327 println!("builder.build(): {:?}", start.elapsed());
328
329 let start = Instant::now();
330 let new_proof = data.prove(pw)?;
331 println!("generate new_proof: {:?}", start.elapsed());
332
333 let start = Instant::now();
334 data.verify(new_proof.clone())?;
335 println!("verify new_proof: {:?}", start.elapsed());
336
337 #[cfg(test)]
338 data.verifier_data().verify(ProofWithPublicInputs {
339 proof: new_proof.proof.clone(),
340 public_inputs: new_proof.public_inputs.clone(),
341 })?;
342
343 #[cfg(test)]
344 verifier_data.verify(ProofWithPublicInputs {
345 proof: new_proof.proof.clone(),
346 public_inputs: new_proof.public_inputs.clone(),
347 })?;
348
349 Ok(new_proof.proof)
350 }
351}
352
353#[cfg(test)]
354mod tests {
355 use anyhow::Result;
356 use hashbrown::HashMap;
357 use plonky2::field::types::Field;
358 use plonky2::plonk::proof::ProofWithPublicInputs;
359 use plonky2::recursion::dummy_circuit::cyclic_base_proof;
360 use rand;
361 use std::time::Instant;
362
363 use super::*;
364
365 // this sets the plonky2 internal logs level
366 fn set_log() {
367 let _ = env_logger::builder()
368 .filter_level(log::LevelFilter::Warn)
369 .is_test(true)
370 .try_init();
371 }
372
373 /// to run:
374 /// cargo test --release test_tree_recursion -- --nocapture
375 #[test]
376 fn test_tree_recursion() -> Result<()> {
377 // For testing: change the following `N` value to try different arities of the recursion tree:
378 test_tree_recursion_opt::<2>()?; // N=2
379
380 test_tree_recursion_opt::<3>()?; // N=3
381
382 Ok(())
383 }
384
385 fn test_tree_recursion_opt<const N: usize>() -> Result<()> {
386 set_log();
387 println!("\n--------------------------------------------------");
388 println!("\n--------------------------------------------------");
389 println!(
390 "\nrunning test:\n test_tree_recursion_opt with N={} (arity)",
391 N
392 );
393
394 let l: u32 = 2; // levels of the recursion (binary) tree
395 let k = (N as u32).pow(l) as usize; // number of leafs in the recursion tree, N^l
396 println!(
397 "Testing a {}-arity recursion tree, of {} levels, with {} leaves",
398 N, l, k
399 );
400
401 let mut rng: rand::rngs::ThreadRng = rand::thread_rng();
402 let schnorr = SchnorrSigner::new();
403 const MSG_LEN: usize = 5;
404 let msg: Vec<F> = schnorr.u64_into_goldilocks_vec(vec![1500, 1600, 2, 2, 2]);
405 assert_eq!(msg.len(), MSG_LEN);
406
407 // generate k key pairs
408 let sk_vec: Vec<SchnorrSecretKey> =
409 (0..k).map(|i| SchnorrSecretKey { sk: i as u64 }).collect();
410 let pk_vec: Vec<SchnorrPublicKey> = sk_vec.iter().map(|&sk| schnorr.keygen(&sk)).collect();
411
412 let sig_vec: Vec<SchnorrSignature> = sk_vec
413 .iter()
414 .map(|&sk| schnorr.sign(&msg, &sk, &mut rng))
415 .collect();
416
417 // build the circuit_data & verifier_data for the recursive circuit
418 let circuit_data = Recursion::<N>::circuit_data(MSG_LEN)?;
419 let verifier_data = circuit_data.verifier_data();
420
421 let dummy_proof_pis = cyclic_base_proof(
422 &circuit_data.common,
423 &verifier_data.verifier_only,
424 HashMap::new(),
425 );
426 let dummy_proof = dummy_proof_pis.proof;
427
428 // we start with k dummy proofs, since at the leafs level we don't have proofs yet and we
429 // just verify the signatures
430 let mut proofs_at_level_i: Vec<PlonkyProof> =
431 (0..k).into_iter().map(|_| dummy_proof.clone()).collect();
432
433 // loop over the recursion levels
434 for i in 0..l {
435 println!("\n=== recursion level i={}", i);
436 let mut next_level_proofs: Vec<PlonkyProof> = vec![];
437
438 // loop over the nodes of each recursion tree level
439 for j in (0..proofs_at_level_i.len()).into_iter().step_by(N) {
440 println!(
441 "\n------ recursion node: (level) i={}, (node in level) j={}",
442 i, j
443 );
444
445 // - if we're at the first level of the recursion tree:
446 // proof_enabled=false=0, so that the circuit verifies the signature and not the proof.
447 // - else:
448 // proof_enabled=true=1, so that the circuit verifies the proof and not the signature.
449 //
450 // In future tests we will try other cases (eg. some sigs and some proofs in a
451 // node), but for the moment we just do base_case: sig verify, other cases: proof
452 // verify.
453 let proof_enabled = if i == 0 { F::ZERO } else { F::ONE };
454
455 // prepare the inputs for the `Recursion::prove_step` call
456 let proofs_selectors = (0..N).into_iter().map(|_| proof_enabled.clone()).collect();
457 let pks = (0..N)
458 .into_iter()
459 .enumerate()
460 .map(|(k, _)| pk_vec[j + k].clone())
461 .collect();
462 let sigs = (0..N)
463 .into_iter()
464 .enumerate()
465 .map(|(k, _)| sig_vec[j + k].clone())
466 .collect();
467 let proofs = (0..N)
468 .into_iter()
469 .enumerate()
470 .map(|(k, _)| proofs_at_level_i[j + k].clone())
471 .collect();
472
473 // do the recursive step
474 let start = Instant::now();
475 let new_proof = Recursion::<N>::prove_step(
476 verifier_data.clone(),
477 &msg,
478 proofs_selectors,
479 &pks,
480 &sigs,
481 &proofs,
482 )?;
483 println!(
484 "Recursion::prove_step (level: i={}, node: j={}) took: {:?}",
485 i,
486 j,
487 start.elapsed()
488 );
489
490 // verify the recursive proof
491 let public_inputs = RecursiveCircuit::<N>::prepare_public_inputs(
492 verifier_data.clone(),
493 msg.clone(),
494 );
495 verifier_data.clone().verify(ProofWithPublicInputs {
496 proof: new_proof.clone(),
497 public_inputs: public_inputs.clone(),
498 })?;
499
500 // set new_proof for next iteration
501 next_level_proofs.push(new_proof);
502 }
503 proofs_at_level_i = next_level_proofs.clone();
504 }
505 assert_eq!(proofs_at_level_i.len(), 1);
506 let last_proof = proofs_at_level_i[0].clone();
507
508 // verify the last proof
509 let public_inputs =
510 RecursiveCircuit::<N>::prepare_public_inputs(verifier_data.clone(), msg.clone());
511 verifier_data.clone().verify(ProofWithPublicInputs {
512 proof: last_proof.clone(),
513 public_inputs: public_inputs.clone(),
514 })?;
515
516 Ok(())
517 }
518}