1pub mod zkr;
16
17use std::{
18 collections::{BTreeMap, VecDeque},
19 fmt::Debug,
20 sync::Mutex,
21};
22
23use anyhow::{anyhow, bail, ensure, Context, Result};
24use risc0_binfmt::read_sha_halfs;
25use risc0_circuit_recursion::{
26 control_id::BN254_IDENTITY_CONTROL_ID,
27 prove::{DigestKind, RecursionReceipt},
28 CircuitImpl,
29};
30use risc0_circuit_rv32im::control_id::POSEIDON2_CONTROL_IDS;
31use risc0_zkp::{
32 adapter::{CircuitInfo, PROOF_SYSTEM_INFO},
33 core::{
34 digest::{Digest, DIGEST_SHORTS},
35 hash::hash_suite_from_name,
36 },
37 field::baby_bear::{BabyBear, BabyBearElem, BabyBearExtElem},
38 hal::{CircuitHal, Hal},
39 verify::ReadIOP,
40 MIN_CYCLES_PO2,
41};
42use serde::Serialize;
43
44use crate::{
45 receipt::{
46 merkle::{MerkleGroup, MerkleProof},
47 SegmentReceipt, SuccinctReceipt, SuccinctReceiptVerifierParameters,
48 },
49 receipt_claim::{Assumption, MaybePruned, Merge},
50 sha::Digestible,
51 ProverOpts, ReceiptClaim, Unknown,
52};
53
54use risc0_circuit_recursion::prove::Program;
55
56pub const RECURSION_PO2: usize = 18;
58
59pub(crate) type ZkrRegistryEntry = Box<dyn Fn() -> Result<Program> + Send + 'static>;
60
61pub(crate) type ZkrRegistry = BTreeMap<Digest, ZkrRegistryEntry>;
62
63pub(crate) static ZKR_REGISTRY: Mutex<ZkrRegistry> = Mutex::new(BTreeMap::new());
65
66pub fn lift(segment_receipt: &SegmentReceipt) -> Result<SuccinctReceipt<ReceiptClaim>> {
73 tracing::debug!("Proving lift: claim = {:#?}", segment_receipt.claim);
74 let opts = ProverOpts::succinct();
75 let mut prover = Prover::new_lift(segment_receipt, opts.clone())?;
76
77 let receipt = prover.prover.run()?;
78 let mut out_stream = VecDeque::<u32>::new();
79 out_stream.extend(receipt.output.iter());
80 let claim_decoded = ReceiptClaim::decode(&mut out_stream)?;
81 tracing::debug!("Proving lift finished: decoded claim = {claim_decoded:#?}");
82
83 let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
85 .get_proof(&prover.control_id, opts.hash_suite()?.hashfn.as_ref())?;
86 Ok(SuccinctReceipt {
87 seal: receipt.seal,
88 hashfn: opts.hashfn,
89 control_id: prover.control_id,
90 control_inclusion_proof,
91 claim: claim_decoded.merge(&segment_receipt.claim)?.into(),
92 verifier_parameters: SuccinctReceiptVerifierParameters::default().digest(),
93 })
94}
95
96pub fn join(
101 a: &SuccinctReceipt<ReceiptClaim>,
102 b: &SuccinctReceipt<ReceiptClaim>,
103) -> Result<SuccinctReceipt<ReceiptClaim>> {
104 tracing::debug!("Proving join: a.claim = {:#?}", a.claim);
105 tracing::debug!("Proving join: b.claim = {:#?}", b.claim);
106
107 let opts = ProverOpts::succinct();
108 let mut prover = Prover::new_join(a, b, opts.clone())?;
109 let receipt = prover.prover.run()?;
110 let mut out_stream = VecDeque::<u32>::new();
111 out_stream.extend(receipt.output.iter());
112
113 let ab_claim = ReceiptClaim {
115 pre: a.claim.as_value()?.pre.clone(),
116 post: b.claim.as_value()?.post.clone(),
117 exit_code: b.claim.as_value()?.exit_code,
118 input: a.claim.as_value()?.input.clone(),
119 output: b.claim.as_value()?.output.clone(),
120 };
121
122 let claim_decoded = ReceiptClaim::decode(&mut out_stream)?;
123 tracing::debug!("Proving join finished: decoded claim = {claim_decoded:#?}");
124
125 let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
127 .get_proof(&prover.control_id, opts.hash_suite()?.hashfn.as_ref())?;
128 Ok(SuccinctReceipt {
129 seal: receipt.seal,
130 hashfn: opts.hashfn,
131 control_id: prover.control_id,
132 control_inclusion_proof,
133 claim: claim_decoded.merge(&ab_claim)?.into(),
134 verifier_parameters: SuccinctReceiptVerifierParameters::default().digest(),
135 })
136}
137
138pub fn resolve<Claim>(
144 conditional: &SuccinctReceipt<ReceiptClaim>,
145 assumption: &SuccinctReceipt<Claim>,
146) -> Result<SuccinctReceipt<ReceiptClaim>>
147where
148 Claim: risc0_binfmt::Digestible + Debug + Clone + Serialize,
149{
150 tracing::debug!(
151 "Proving resolve: conditional.claim = {:#?}",
152 conditional.claim,
153 );
154 tracing::debug!(
155 "Proving resolve: assumption.claim = {:#?}",
156 assumption.claim,
157 );
158
159 let mut resolved_claim = conditional
162 .claim
163 .as_value()
164 .context("conditional receipt claim is pruned")?
165 .clone();
166 resolved_claim
170 .output
171 .as_value_mut()
172 .context("conditional receipt output is pruned")?
173 .as_mut()
174 .ok_or(anyhow!(
175 "conditional receipt has empty output and no assumptions"
176 ))?
177 .assumptions
178 .as_value_mut()
179 .context("conditional receipt assumptions are pruned")?
180 .0
181 .drain(..1)
182 .next()
183 .ok_or(anyhow!(
184 "cannot resolve assumption from receipt with no assumptions"
185 ))?;
186
187 let opts = ProverOpts::succinct();
188 let mut prover = Prover::new_resolve(conditional, assumption, opts.clone())?;
189 let receipt = prover.prover.run()?;
190 let mut out_stream = VecDeque::<u32>::new();
191 out_stream.extend(receipt.output.iter());
192
193 let claim_decoded = ReceiptClaim::decode(&mut out_stream)?;
194 tracing::debug!("Proving resolve finished: decoded claim = {claim_decoded:#?}");
195
196 let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
198 .get_proof(&prover.control_id, opts.hash_suite()?.hashfn.as_ref())?;
199 Ok(SuccinctReceipt {
200 seal: receipt.seal,
201 hashfn: opts.hashfn,
202 control_id: prover.control_id,
203 control_inclusion_proof,
204 claim: claim_decoded.merge(&resolved_claim)?.into(),
205 verifier_parameters: SuccinctReceiptVerifierParameters::default().digest(),
206 })
207}
208
209pub fn identity_p254(a: &SuccinctReceipt<ReceiptClaim>) -> Result<SuccinctReceipt<ReceiptClaim>> {
215 let opts = ProverOpts::succinct()
216 .with_hashfn("poseidon_254".to_string())
217 .with_control_ids(vec![BN254_IDENTITY_CONTROL_ID]);
218
219 let mut prover = Prover::new_identity(a, opts.clone())?;
220 let receipt = prover.prover.run()?;
221 let mut out_stream = VecDeque::<u32>::new();
222 out_stream.extend(receipt.output.iter());
223 let claim = MaybePruned::Value(ReceiptClaim::decode(&mut out_stream)?).merge(&a.claim)?;
224
225 let hashfn = opts.hash_suite()?.hashfn;
227 let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
228 .get_proof(&prover.control_id, hashfn.as_ref())?;
229 let control_root = control_inclusion_proof.root(&prover.control_id, hashfn.as_ref());
230 let params = SuccinctReceiptVerifierParameters {
231 control_root,
232 inner_control_root: Some(a.control_root()?),
233 proof_system_info: PROOF_SYSTEM_INFO,
234 circuit_info: CircuitImpl::CIRCUIT_INFO,
235 };
236 Ok(SuccinctReceipt {
237 seal: receipt.seal,
238 hashfn: opts.hashfn,
239 control_id: prover.control_id,
240 control_inclusion_proof,
241 claim,
242 verifier_parameters: params.digest(),
243 })
244}
245
246pub fn prove_zkr(
248 program: Program,
249 control_id: &Digest,
250 allowed_control_ids: Vec<Digest>,
251 input: &[u8],
252) -> Result<SuccinctReceipt<Unknown>> {
253 let opts = ProverOpts::succinct().with_control_ids(allowed_control_ids);
254 let mut prover = Prover::new(program, *control_id, opts.clone());
255 prover.add_input(bytemuck::cast_slice(input));
256
257 tracing::debug!("Running prover");
258 let receipt = prover.run()?;
259
260 tracing::trace!("zkr receipt: {receipt:?}");
261
262 let claim_digest: Digest = read_sha_halfs(&mut VecDeque::from_iter(
264 bytemuck::checked::cast_slice::<_, BabyBearElem>(
265 &receipt.seal[DIGEST_SHORTS..2 * DIGEST_SHORTS],
266 )
267 .iter()
268 .copied()
269 .map(u32::from),
270 ))?;
271
272 let hashfn = opts.hash_suite()?.hashfn;
273 let control_inclusion_proof =
274 MerkleGroup::new(opts.control_ids.clone())?.get_proof(control_id, hashfn.as_ref())?;
275
276 Ok(SuccinctReceipt {
277 seal: receipt.seal,
278 hashfn: opts.hashfn,
279 control_id: *control_id,
280 control_inclusion_proof,
281 claim: MaybePruned::<Unknown>::Pruned(claim_digest),
282 verifier_parameters: SuccinctReceiptVerifierParameters::default().digest(),
283 })
284}
285
286pub fn prove_registered_zkr(
288 control_id: &Digest,
289 allowed_control_ids: Vec<Digest>,
290 input: &[u8],
291) -> Result<SuccinctReceipt<Unknown>> {
292 let zkr = get_registered_zkr(control_id)?;
293 prove_zkr(zkr, control_id, allowed_control_ids, input)
294}
295
296pub fn register_zkr(
298 control_id: &Digest,
299 get_program_fn: impl Fn() -> Result<Program> + Send + 'static,
300) -> Option<ZkrRegistryEntry> {
301 let mut registry = ZKR_REGISTRY.lock().unwrap();
302 registry.insert(*control_id, Box::new(get_program_fn))
303}
304
305pub fn get_registered_zkr(control_id: &Digest) -> Result<Program> {
307 let registry = ZKR_REGISTRY.lock().unwrap();
308 registry
309 .get(control_id)
310 .map(|f| f())
311 .unwrap_or_else(|| bail!("Control id {control_id} unregistered"))
312}
313
314#[cfg(test)]
319pub fn test_zkr(
320 digest1: &Digest,
321 digest2: &Digest,
322 po2: usize,
323) -> Result<SuccinctReceipt<crate::receipt_claim::Unknown>> {
324 use risc0_circuit_recursion::prove::zkr::get_zkr;
325 use risc0_zkp::core::hash::poseidon2::Poseidon2HashSuite;
326
327 let program = get_zkr("test_recursion_circuit.zkr", po2)?;
328 let suite = Poseidon2HashSuite::new_suite();
329 let control_id = program.compute_control_id(suite.clone());
330 let opts = ProverOpts::succinct().with_control_ids(vec![control_id]);
331
332 let mut prover = Prover::new(program, control_id, opts.clone());
333 prover.add_input_digest(digest1, DigestKind::Poseidon2);
334 prover.add_input_digest(digest2, DigestKind::Poseidon2);
335
336 let receipt = prover.run()?;
337
338 let claim_digest = risc0_binfmt::read_sha_halfs(&mut VecDeque::from_iter(
340 bytemuck::checked::cast_slice::<_, BabyBearElem>(
341 &receipt.seal[DIGEST_SHORTS..2 * DIGEST_SHORTS],
342 )
343 .iter()
344 .copied()
345 .map(u32::from),
346 ))?;
347
348 let hashfn = opts.hash_suite()?.hashfn;
350 let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
351 .get_proof(&prover.control_id, hashfn.as_ref())?;
352 let control_root = control_inclusion_proof.root(&prover.control_id, hashfn.as_ref());
353 let params = SuccinctReceiptVerifierParameters {
354 control_root,
355 inner_control_root: Some(digest1.to_owned()),
356 proof_system_info: PROOF_SYSTEM_INFO,
357 circuit_info: CircuitImpl::CIRCUIT_INFO,
358 };
359 Ok(SuccinctReceipt {
360 seal: receipt.seal,
361 hashfn: suite.name,
362 control_id: prover.control_id,
363 control_inclusion_proof,
364 claim: MaybePruned::Pruned(claim_digest),
365 verifier_parameters: params.digest(),
366 })
367}
368
369pub struct Prover {
371 prover: risc0_circuit_recursion::prove::Prover,
372 control_id: Digest,
373}
374
375impl Prover {
376 pub(crate) fn new(program: Program, control_id: Digest, opts: ProverOpts) -> Self {
377 Self {
378 prover: risc0_circuit_recursion::prove::Prover::new(program, &opts.hashfn),
379 control_id,
380 }
381 }
382
383 pub fn control_id(&self) -> &Digest {
385 &self.control_id
386 }
387
388 pub fn new_test_recursion_circuit(digests: [&Digest; 2], opts: ProverOpts) -> Result<Self> {
391 let (program, control_id) = zkr::test_recursion_circuit(&opts.hashfn)?;
392 let mut prover = Prover::new(program, control_id, opts);
393
394 for digest in digests {
395 prover.add_input_digest(digest, DigestKind::Poseidon2);
396 }
397
398 Ok(prover)
399 }
400
401 pub fn new_lift(segment: &SegmentReceipt, opts: ProverOpts) -> Result<Self> {
410 ensure!(
411 segment.hashfn == "poseidon2",
412 "lift recursion program only supports poseidon2 hashfn; received {}",
413 segment.hashfn
414 );
415
416 let inner_hash_suite = hash_suite_from_name(&segment.hashfn)
417 .ok_or_else(|| anyhow!("unsupported hash function: {}", segment.hashfn))?;
418 let allowed_ids = MerkleGroup::new(opts.control_ids.clone())?;
419 let merkle_root = allowed_ids.calc_root(inner_hash_suite.hashfn.as_ref());
420
421 let mut iop = ReadIOP::new(&segment.seal, inner_hash_suite.rng.as_ref());
424 iop.read_field_elem_slice::<BabyBearElem>(risc0_circuit_rv32im::CircuitImpl::OUTPUT_SIZE);
425 let po2 = *iop.read_u32s(1).first().unwrap() as usize;
426
427 let (program, control_id) = zkr::lift(po2, &opts.hashfn)?;
429 let mut prover = Prover::new(program, control_id, opts);
430
431 prover.add_input_digest(&merkle_root, DigestKind::Poseidon2);
432
433 let which = po2 - MIN_CYCLES_PO2;
435 let inner_control_id = POSEIDON2_CONTROL_IDS[which];
436 prover.add_seal(
437 &segment.seal,
438 &inner_control_id,
439 &allowed_ids.get_proof(&inner_control_id, inner_hash_suite.hashfn.as_ref())?,
440 )?;
441
442 Ok(prover)
443 }
444
445 pub fn new_join(
451 a: &SuccinctReceipt<ReceiptClaim>,
452 b: &SuccinctReceipt<ReceiptClaim>,
453 opts: ProverOpts,
454 ) -> Result<Self> {
455 ensure!(
456 a.hashfn == "poseidon2",
457 "join recursion program only supports poseidon2 hashfn; received {}",
458 a.hashfn
459 );
460 ensure!(
461 b.hashfn == "poseidon2",
462 "join recursion program only supports poseidon2 hashfn; received {}",
463 b.hashfn
464 );
465
466 let (program, control_id) = zkr::join(&opts.hashfn)?;
467 let mut prover = Prover::new(program, control_id, opts);
468
469 let merkle_root = a.control_root()?;
473 ensure!(
474 merkle_root == b.control_root()?,
475 "merkle roots for a and b do not match: {} != {}",
476 merkle_root,
477 b.control_root()?
478 );
479
480 prover.add_input_digest(&merkle_root, DigestKind::Poseidon2);
481 prover.add_segment_receipt(a)?;
482 prover.add_segment_receipt(b)?;
483 Ok(prover)
484 }
485
486 pub fn new_resolve<Claim>(
493 cond: &SuccinctReceipt<ReceiptClaim>,
494 assum: &SuccinctReceipt<Claim>,
495 opts: ProverOpts,
496 ) -> Result<Self>
497 where
498 Claim: risc0_binfmt::Digestible + Debug + Clone + Serialize,
499 {
500 ensure!(
501 cond.hashfn == "poseidon2",
502 "resolve recursion program only supports poseidon2 hashfn; received {}",
503 cond.hashfn
504 );
505 ensure!(
506 assum.hashfn == "poseidon2",
507 "resolve recursion program only supports poseidon2 hashfn; received {}",
508 assum.hashfn
509 );
510
511 let (program, control_id) = zkr::resolve(&opts.hashfn)?;
513 let mut prover = Prover::new(program, control_id, opts);
514
515 prover.add_input_digest(&cond.control_root()?, DigestKind::Poseidon2);
519 prover.add_segment_receipt(cond)?;
520
521 let output = cond
522 .claim
523 .as_value()
524 .context("cannot resolve conditional receipt with pruned claim")?
525 .output
526 .as_value()
527 .context("cannot resolve conditional receipt with pruned output")?
528 .as_ref()
529 .ok_or(anyhow!("cannot resolve conditional receipt with no output"))?
530 .clone();
531
532 let assumptions = output
535 .assumptions
536 .value()
537 .context("cannot resolve conditional receipt with pruned assumptions")?;
538 let head: Assumption = assumptions
539 .0
540 .first()
541 .ok_or(anyhow!(
542 "cannot resolve conditional receipt with no assumptions"
543 ))?
544 .as_value()
545 .context("cannot resolve conditional receipt with pruned head assumption")?
546 .clone();
547
548 ensure!(
550 head.claim == assum.claim.digest(),
551 "assumption receipt claim does not match head of assumptions list"
552 );
553 let expected_root = match head.control_root == Digest::ZERO {
554 true => cond.control_root()?,
555 false => head.control_root,
556 };
557 ensure!(
558 expected_root == assum.control_root()?,
559 "assumption receipt control root does not match head of assumptions list"
560 );
561
562 let mut assumptions_tail = assumptions;
563 assumptions_tail.resolve(&head.digest())?;
564
565 prover.add_assumption_receipt(head, assum)?;
566 prover.add_input_digest(&assumptions_tail.digest(), DigestKind::Sha256);
567 prover.add_input_digest(&output.journal.digest(), DigestKind::Sha256);
568 Ok(prover)
569 }
570
571 pub fn new_identity(a: &SuccinctReceipt<ReceiptClaim>, opts: ProverOpts) -> Result<Self> {
576 ensure!(
577 a.hashfn == "poseidon2",
578 "identity recursion program only supports poseidon2 hashfn; received {}",
579 a.hashfn
580 );
581
582 let (program, control_id) = zkr::identity(&opts.hashfn)?;
583 let mut prover = Prover::new(program, control_id, opts);
584
585 prover.add_input_digest(&a.control_root()?, DigestKind::Poseidon2);
586 prover.add_segment_receipt(a)?;
587 Ok(prover)
588 }
589
590 pub(crate) fn add_input(&mut self, input: &[u32]) {
591 self.prover.add_input(input)
592 }
593
594 fn add_input_digest(&mut self, digest: &Digest, kind: DigestKind) {
596 self.prover.add_input_digest(digest, kind)
597 }
598
599 pub fn add_seal(
601 &mut self,
602 seal: &[u32],
603 control_id: &Digest,
604 control_inclusion_proof: &MerkleProof,
605 ) -> Result<()> {
606 tracing::debug!("Control ID = {:?}", control_id);
607 self.add_input(seal);
608 tracing::debug!("index = {:?}", control_inclusion_proof.index);
609 self.add_input(bytemuck::cast_slice(&[BabyBearElem::new(
610 control_inclusion_proof.index,
611 )]));
612 for digest in &control_inclusion_proof.digests {
613 tracing::debug!("path = {:?}", digest);
614 self.add_input_digest(digest, DigestKind::Poseidon2);
615 }
616 Ok(())
617 }
618
619 fn add_assumption_receipt<Claim>(
621 &mut self,
622 assumption: Assumption,
623 receipt: &SuccinctReceipt<Claim>,
624 ) -> Result<()>
625 where
626 Claim: risc0_binfmt::Digestible + Debug + Clone + Serialize,
627 {
628 self.add_seal(
629 &receipt.seal,
630 &receipt.control_id,
631 &receipt.control_inclusion_proof,
632 )?;
633 let zero_root = BabyBearElem::new((assumption.control_root == Digest::ZERO) as u32);
635 self.add_input(bytemuck::cast_slice(&[zero_root]));
636 Ok(())
637 }
638
639 fn add_segment_receipt(&mut self, a: &SuccinctReceipt<ReceiptClaim>) -> Result<()> {
641 self.add_seal(&a.seal, &a.control_id, &a.control_inclusion_proof)?;
642 let mut data = Vec::<u32>::new();
643 a.claim.as_value()?.encode(&mut data)?;
644 let data_fp: Vec<BabyBearElem> = data.iter().map(|x| BabyBearElem::new(*x)).collect();
645 self.add_input(bytemuck::cast_slice(&data_fp));
646 Ok(())
647 }
648
649 pub fn run(&mut self) -> Result<RecursionReceipt> {
652 self.prover.run()
653 }
654
655 pub fn run_with_hal<H, C>(&mut self, hal: &H, circuit_hal: &C) -> Result<RecursionReceipt>
658 where
659 H: Hal<Field = BabyBear, Elem = BabyBearElem, ExtElem = BabyBearExtElem>,
660 C: CircuitHal<H>,
661 {
662 self.prover.run_with_hal(hal, circuit_hal)
663 }
664}