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