pub mod zkr;
use std::{
collections::{BTreeMap, VecDeque},
fmt::Debug,
sync::Mutex,
};
use anyhow::{anyhow, bail, ensure, Context, Result};
use risc0_binfmt::read_sha_halfs;
use risc0_circuit_recursion::{
control_id::BN254_IDENTITY_CONTROL_ID,
prove::{DigestKind, RecursionReceipt},
CircuitImpl,
};
use risc0_circuit_rv32im::control_id::POSEIDON2_CONTROL_IDS;
use risc0_zkp::{
adapter::{CircuitInfo, PROOF_SYSTEM_INFO},
core::{
digest::{Digest, DIGEST_SHORTS},
hash::hash_suite_from_name,
},
field::baby_bear::{BabyBear, BabyBearElem, BabyBearExtElem},
hal::{CircuitHal, Hal},
verify::ReadIOP,
MIN_CYCLES_PO2,
};
use serde::Serialize;
use crate::{
receipt::{
merkle::{MerkleGroup, MerkleProof},
SegmentReceipt, SuccinctReceipt, SuccinctReceiptVerifierParameters,
},
receipt_claim::{Assumption, MaybePruned, Merge},
sha::Digestible,
ProverOpts, ReceiptClaim, Unknown,
};
use risc0_circuit_recursion::prove::Program;
pub const RECURSION_PO2: usize = 18;
pub(crate) type ZkrRegistryEntry = Box<dyn Fn() -> Result<Program> + Send + 'static>;
pub(crate) type ZkrRegistry = BTreeMap<Digest, ZkrRegistryEntry>;
pub(crate) static ZKR_REGISTRY: Mutex<ZkrRegistry> = Mutex::new(BTreeMap::new());
pub fn lift(segment_receipt: &SegmentReceipt) -> Result<SuccinctReceipt<ReceiptClaim>> {
tracing::debug!("Proving lift: claim = {:#?}", segment_receipt.claim);
let opts = ProverOpts::succinct();
let mut prover = Prover::new_lift(segment_receipt, opts.clone())?;
let receipt = prover.prover.run()?;
let mut out_stream = VecDeque::<u32>::new();
out_stream.extend(receipt.output.iter());
let claim_decoded = ReceiptClaim::decode(&mut out_stream)?;
tracing::debug!("Proving lift finished: decoded claim = {claim_decoded:#?}");
let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
.get_proof(&prover.control_id, opts.hash_suite()?.hashfn.as_ref())?;
Ok(SuccinctReceipt {
seal: receipt.seal,
hashfn: opts.hashfn,
control_id: prover.control_id,
control_inclusion_proof,
claim: claim_decoded.merge(&segment_receipt.claim)?.into(),
verifier_parameters: SuccinctReceiptVerifierParameters::default().digest(),
})
}
pub fn join(
a: &SuccinctReceipt<ReceiptClaim>,
b: &SuccinctReceipt<ReceiptClaim>,
) -> Result<SuccinctReceipt<ReceiptClaim>> {
tracing::debug!("Proving join: a.claim = {:#?}", a.claim);
tracing::debug!("Proving join: b.claim = {:#?}", b.claim);
let opts = ProverOpts::succinct();
let mut prover = Prover::new_join(a, b, opts.clone())?;
let receipt = prover.prover.run()?;
let mut out_stream = VecDeque::<u32>::new();
out_stream.extend(receipt.output.iter());
let ab_claim = ReceiptClaim {
pre: a.claim.as_value()?.pre.clone(),
post: b.claim.as_value()?.post.clone(),
exit_code: b.claim.as_value()?.exit_code,
input: a.claim.as_value()?.input.clone(),
output: b.claim.as_value()?.output.clone(),
};
let claim_decoded = ReceiptClaim::decode(&mut out_stream)?;
tracing::debug!("Proving join finished: decoded claim = {claim_decoded:#?}");
let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
.get_proof(&prover.control_id, opts.hash_suite()?.hashfn.as_ref())?;
Ok(SuccinctReceipt {
seal: receipt.seal,
hashfn: opts.hashfn,
control_id: prover.control_id,
control_inclusion_proof,
claim: claim_decoded.merge(&ab_claim)?.into(),
verifier_parameters: SuccinctReceiptVerifierParameters::default().digest(),
})
}
pub fn resolve<Claim>(
conditional: &SuccinctReceipt<ReceiptClaim>,
assumption: &SuccinctReceipt<Claim>,
) -> Result<SuccinctReceipt<ReceiptClaim>>
where
Claim: risc0_binfmt::Digestible + Debug + Clone + Serialize,
{
tracing::debug!(
"Proving resolve: conditional.claim = {:#?}",
conditional.claim,
);
tracing::debug!(
"Proving resolve: assumption.claim = {:#?}",
assumption.claim,
);
let mut resolved_claim = conditional
.claim
.as_value()
.context("conditional receipt claim is pruned")?
.clone();
resolved_claim
.output
.as_value_mut()
.context("conditional receipt output is pruned")?
.as_mut()
.ok_or(anyhow!(
"conditional receipt has empty output and no assumptions"
))?
.assumptions
.as_value_mut()
.context("conditional receipt assumptions are pruned")?
.0
.drain(..1)
.next()
.ok_or(anyhow!(
"cannot resolve assumption from receipt with no assumptions"
))?;
let opts = ProverOpts::succinct();
let mut prover = Prover::new_resolve(conditional, assumption, opts.clone())?;
let receipt = prover.prover.run()?;
let mut out_stream = VecDeque::<u32>::new();
out_stream.extend(receipt.output.iter());
let claim_decoded = ReceiptClaim::decode(&mut out_stream)?;
tracing::debug!("Proving resolve finished: decoded claim = {claim_decoded:#?}");
let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
.get_proof(&prover.control_id, opts.hash_suite()?.hashfn.as_ref())?;
Ok(SuccinctReceipt {
seal: receipt.seal,
hashfn: opts.hashfn,
control_id: prover.control_id,
control_inclusion_proof,
claim: claim_decoded.merge(&resolved_claim)?.into(),
verifier_parameters: SuccinctReceiptVerifierParameters::default().digest(),
})
}
pub fn identity_p254(a: &SuccinctReceipt<ReceiptClaim>) -> Result<SuccinctReceipt<ReceiptClaim>> {
let opts = ProverOpts::succinct()
.with_hashfn("poseidon_254".to_string())
.with_control_ids(vec![BN254_IDENTITY_CONTROL_ID]);
let mut prover = Prover::new_identity(a, opts.clone())?;
let receipt = prover.prover.run()?;
let mut out_stream = VecDeque::<u32>::new();
out_stream.extend(receipt.output.iter());
let claim = MaybePruned::Value(ReceiptClaim::decode(&mut out_stream)?).merge(&a.claim)?;
let hashfn = opts.hash_suite()?.hashfn;
let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
.get_proof(&prover.control_id, hashfn.as_ref())?;
let control_root = control_inclusion_proof.root(&prover.control_id, hashfn.as_ref());
let params = SuccinctReceiptVerifierParameters {
control_root,
inner_control_root: Some(a.control_root()?),
proof_system_info: PROOF_SYSTEM_INFO,
circuit_info: CircuitImpl::CIRCUIT_INFO,
};
Ok(SuccinctReceipt {
seal: receipt.seal,
hashfn: opts.hashfn,
control_id: prover.control_id,
control_inclusion_proof,
claim,
verifier_parameters: params.digest(),
})
}
pub fn prove_zkr(
program: Program,
control_id: &Digest,
allowed_control_ids: Vec<Digest>,
input: &[u8],
) -> Result<SuccinctReceipt<Unknown>> {
let opts = ProverOpts::succinct().with_control_ids(allowed_control_ids);
let mut prover = Prover::new(program, *control_id, opts.clone());
prover.add_input(bytemuck::cast_slice(input));
tracing::debug!("Running prover");
let receipt = prover.run()?;
tracing::trace!("zkr receipt: {receipt:?}");
let claim_digest: Digest = read_sha_halfs(&mut VecDeque::from_iter(
bytemuck::checked::cast_slice::<_, BabyBearElem>(
&receipt.seal[DIGEST_SHORTS..2 * DIGEST_SHORTS],
)
.iter()
.copied()
.map(u32::from),
))?;
let hashfn = opts.hash_suite()?.hashfn;
let control_inclusion_proof =
MerkleGroup::new(opts.control_ids.clone())?.get_proof(control_id, hashfn.as_ref())?;
Ok(SuccinctReceipt {
seal: receipt.seal,
hashfn: opts.hashfn,
control_id: *control_id,
control_inclusion_proof,
claim: MaybePruned::<Unknown>::Pruned(claim_digest),
verifier_parameters: SuccinctReceiptVerifierParameters::default().digest(),
})
}
pub fn prove_registered_zkr(
control_id: &Digest,
allowed_control_ids: Vec<Digest>,
input: &[u8],
) -> Result<SuccinctReceipt<Unknown>> {
let zkr = get_registered_zkr(control_id)?;
prove_zkr(zkr, control_id, allowed_control_ids, input)
}
pub fn register_zkr(
control_id: &Digest,
get_program_fn: impl Fn() -> Result<Program> + Send + 'static,
) -> Option<ZkrRegistryEntry> {
let mut registry = ZKR_REGISTRY.lock().unwrap();
registry.insert(*control_id, Box::new(get_program_fn))
}
pub fn get_registered_zkr(control_id: &Digest) -> Result<Program> {
let registry = ZKR_REGISTRY.lock().unwrap();
registry
.get(control_id)
.map(|f| f())
.unwrap_or_else(|| bail!("Control id {control_id} unregistered"))
}
#[cfg(test)]
pub fn test_zkr(
digest1: &Digest,
digest2: &Digest,
po2: usize,
) -> Result<SuccinctReceipt<crate::receipt_claim::Unknown>> {
use risc0_circuit_recursion::prove::zkr::get_zkr;
use risc0_zkp::core::hash::poseidon2::Poseidon2HashSuite;
let program = get_zkr("test_recursion_circuit.zkr", po2)?;
let suite = Poseidon2HashSuite::new_suite();
let control_id = program.compute_control_id(suite.clone());
let opts = ProverOpts::succinct().with_control_ids(vec![control_id]);
let mut prover = Prover::new(program, control_id, opts.clone());
prover.add_input_digest(digest1, DigestKind::Poseidon2);
prover.add_input_digest(digest2, DigestKind::Poseidon2);
let receipt = prover.run()?;
let claim_digest = risc0_binfmt::read_sha_halfs(&mut VecDeque::from_iter(
bytemuck::checked::cast_slice::<_, BabyBearElem>(
&receipt.seal[DIGEST_SHORTS..2 * DIGEST_SHORTS],
)
.iter()
.copied()
.map(u32::from),
))?;
let hashfn = opts.hash_suite()?.hashfn;
let control_inclusion_proof = MerkleGroup::new(opts.control_ids.clone())?
.get_proof(&prover.control_id, hashfn.as_ref())?;
let control_root = control_inclusion_proof.root(&prover.control_id, hashfn.as_ref());
let params = SuccinctReceiptVerifierParameters {
control_root,
inner_control_root: Some(digest1.to_owned()),
proof_system_info: PROOF_SYSTEM_INFO,
circuit_info: CircuitImpl::CIRCUIT_INFO,
};
Ok(SuccinctReceipt {
seal: receipt.seal,
hashfn: suite.name,
control_id: prover.control_id,
control_inclusion_proof,
claim: MaybePruned::Pruned(claim_digest),
verifier_parameters: params.digest(),
})
}
pub struct Prover {
prover: risc0_circuit_recursion::prove::Prover,
control_id: Digest,
}
impl Prover {
pub(crate) fn new(program: Program, control_id: Digest, opts: ProverOpts) -> Self {
Self {
prover: risc0_circuit_recursion::prove::Prover::new(program, &opts.hashfn),
control_id,
}
}
pub fn control_id(&self) -> &Digest {
&self.control_id
}
pub fn new_test_recursion_circuit(digests: [&Digest; 2], opts: ProverOpts) -> Result<Self> {
let (program, control_id) = zkr::test_recursion_circuit(&opts.hashfn)?;
let mut prover = Prover::new(program, control_id, opts);
for digest in digests {
prover.add_input_digest(digest, DigestKind::Poseidon2);
}
Ok(prover)
}
pub fn new_lift(segment: &SegmentReceipt, opts: ProverOpts) -> Result<Self> {
ensure!(
segment.hashfn == "poseidon2",
"lift recursion program only supports poseidon2 hashfn; received {}",
segment.hashfn
);
let inner_hash_suite = hash_suite_from_name(&segment.hashfn)
.ok_or_else(|| anyhow!("unsupported hash function: {}", segment.hashfn))?;
let allowed_ids = MerkleGroup::new(opts.control_ids.clone())?;
let merkle_root = allowed_ids.calc_root(inner_hash_suite.hashfn.as_ref());
let mut iop = ReadIOP::new(&segment.seal, inner_hash_suite.rng.as_ref());
iop.read_field_elem_slice::<BabyBearElem>(risc0_circuit_rv32im::CircuitImpl::OUTPUT_SIZE);
let po2 = *iop.read_u32s(1).first().unwrap() as usize;
let (program, control_id) = zkr::lift(po2, &opts.hashfn)?;
let mut prover = Prover::new(program, control_id, opts);
prover.add_input_digest(&merkle_root, DigestKind::Poseidon2);
let which = po2 - MIN_CYCLES_PO2;
let inner_control_id = POSEIDON2_CONTROL_IDS[which];
prover.add_seal(
&segment.seal,
&inner_control_id,
&allowed_ids.get_proof(&inner_control_id, inner_hash_suite.hashfn.as_ref())?,
)?;
Ok(prover)
}
pub fn new_join(
a: &SuccinctReceipt<ReceiptClaim>,
b: &SuccinctReceipt<ReceiptClaim>,
opts: ProverOpts,
) -> Result<Self> {
ensure!(
a.hashfn == "poseidon2",
"join recursion program only supports poseidon2 hashfn; received {}",
a.hashfn
);
ensure!(
b.hashfn == "poseidon2",
"join recursion program only supports poseidon2 hashfn; received {}",
b.hashfn
);
let (program, control_id) = zkr::join(&opts.hashfn)?;
let mut prover = Prover::new(program, control_id, opts);
let merkle_root = a.control_root()?;
ensure!(
merkle_root == b.control_root()?,
"merkle roots for a and b do not match: {} != {}",
merkle_root,
b.control_root()?
);
prover.add_input_digest(&merkle_root, DigestKind::Poseidon2);
prover.add_segment_receipt(a)?;
prover.add_segment_receipt(b)?;
Ok(prover)
}
pub fn new_resolve<Claim>(
cond: &SuccinctReceipt<ReceiptClaim>,
assum: &SuccinctReceipt<Claim>,
opts: ProverOpts,
) -> Result<Self>
where
Claim: risc0_binfmt::Digestible + Debug + Clone + Serialize,
{
ensure!(
cond.hashfn == "poseidon2",
"resolve recursion program only supports poseidon2 hashfn; received {}",
cond.hashfn
);
ensure!(
assum.hashfn == "poseidon2",
"resolve recursion program only supports poseidon2 hashfn; received {}",
assum.hashfn
);
let (program, control_id) = zkr::resolve(&opts.hashfn)?;
let mut prover = Prover::new(program, control_id, opts);
prover.add_input_digest(&cond.control_root()?, DigestKind::Poseidon2);
prover.add_segment_receipt(cond)?;
let output = cond
.claim
.as_value()
.context("cannot resolve conditional receipt with pruned claim")?
.output
.as_value()
.context("cannot resolve conditional receipt with pruned output")?
.as_ref()
.ok_or(anyhow!("cannot resolve conditional receipt with no output"))?
.clone();
let assumptions = output
.assumptions
.value()
.context("cannot resolve conditional receipt with pruned assumptions")?;
let head: Assumption = assumptions
.0
.first()
.ok_or(anyhow!(
"cannot resolve conditional receipt with no assumptions"
))?
.as_value()
.context("cannot resolve conditional receipt with pruned head assumption")?
.clone();
ensure!(
head.claim == assum.claim.digest(),
"assumption receipt claim does not match head of assumptions list"
);
let expected_root = match head.control_root == Digest::ZERO {
true => cond.control_root()?,
false => head.control_root,
};
ensure!(
expected_root == assum.control_root()?,
"assumption receipt control root does not match head of assumptions list"
);
let mut assumptions_tail = assumptions;
assumptions_tail.resolve(&head.digest())?;
prover.add_assumption_receipt(head, assum)?;
prover.add_input_digest(&assumptions_tail.digest(), DigestKind::Sha256);
prover.add_input_digest(&output.journal.digest(), DigestKind::Sha256);
Ok(prover)
}
pub fn new_identity(a: &SuccinctReceipt<ReceiptClaim>, opts: ProverOpts) -> Result<Self> {
ensure!(
a.hashfn == "poseidon2",
"identity recursion program only supports poseidon2 hashfn; received {}",
a.hashfn
);
let (program, control_id) = zkr::identity(&opts.hashfn)?;
let mut prover = Prover::new(program, control_id, opts);
prover.add_input_digest(&a.control_root()?, DigestKind::Poseidon2);
prover.add_segment_receipt(a)?;
Ok(prover)
}
pub(crate) fn add_input(&mut self, input: &[u32]) {
self.prover.add_input(input)
}
fn add_input_digest(&mut self, digest: &Digest, kind: DigestKind) {
self.prover.add_input_digest(digest, kind)
}
pub fn add_seal(
&mut self,
seal: &[u32],
control_id: &Digest,
control_inclusion_proof: &MerkleProof,
) -> Result<()> {
tracing::debug!("Control ID = {:?}", control_id);
self.add_input(seal);
tracing::debug!("index = {:?}", control_inclusion_proof.index);
self.add_input(bytemuck::cast_slice(&[BabyBearElem::new(
control_inclusion_proof.index,
)]));
for digest in &control_inclusion_proof.digests {
tracing::debug!("path = {:?}", digest);
self.add_input_digest(digest, DigestKind::Poseidon2);
}
Ok(())
}
fn add_assumption_receipt<Claim>(
&mut self,
assumption: Assumption,
receipt: &SuccinctReceipt<Claim>,
) -> Result<()>
where
Claim: risc0_binfmt::Digestible + Debug + Clone + Serialize,
{
self.add_seal(
&receipt.seal,
&receipt.control_id,
&receipt.control_inclusion_proof,
)?;
let zero_root = BabyBearElem::new((assumption.control_root == Digest::ZERO) as u32);
self.add_input(bytemuck::cast_slice(&[zero_root]));
Ok(())
}
fn add_segment_receipt(&mut self, a: &SuccinctReceipt<ReceiptClaim>) -> Result<()> {
self.add_seal(&a.seal, &a.control_id, &a.control_inclusion_proof)?;
let mut data = Vec::<u32>::new();
a.claim.as_value()?.encode(&mut data)?;
let data_fp: Vec<BabyBearElem> = data.iter().map(|x| BabyBearElem::new(*x)).collect();
self.add_input(bytemuck::cast_slice(&data_fp));
Ok(())
}
pub fn run(&mut self) -> Result<RecursionReceipt> {
self.prover.run()
}
pub fn run_with_hal<H, C>(&mut self, hal: &H, circuit_hal: &C) -> Result<RecursionReceipt>
where
H: Hal<Field = BabyBear, Elem = BabyBearElem, ExtElem = BabyBearExtElem>,
C: CircuitHal<H>,
{
self.prover.run_with_hal(hal, circuit_hal)
}
}