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