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