openvm_continuations/prover/root/
trace.rs

1use itertools::Itertools;
2use openvm_circuit::system::memory::merkle::public_values::UserPublicValuesProof;
3use openvm_recursion_circuit::system::{
4    AggregationSubCircuit, CachedTraceCtx, VerifierExternalData, VerifierTraceGen,
5};
6use openvm_stark_backend::{
7    proof::Proof,
8    prover::{ProverBackend, ProvingContext},
9};
10use openvm_stark_sdk::config::baby_bear_poseidon2::{
11    default_duplex_sponge_recorder, DIGEST_SIZE, EF, F,
12};
13use tracing::instrument;
14
15use super::RootProver;
16use crate::{
17    circuit::{deferral::DeferralMerkleProofs, root::RootTraceGen},
18    RootSC, SC,
19};
20
21impl<S: AggregationSubCircuit, T> RootProver<S, T> {
22    pub fn generate_proving_ctx<PB, DC: Clone + Send + Sync>(
23        &self,
24        proof: Proof<SC>,
25        user_pvs_proof: &UserPublicValuesProof<DIGEST_SIZE, PB::Val>,
26        deferral_merkle_proofs: Option<&DeferralMerkleProofs<PB::Val>>,
27        device_ctx: &DC,
28    ) -> Option<ProvingContext<PB>>
29    where
30        PB: ProverBackend<Val = F, Challenge = EF>,
31        PB::Matrix: Clone,
32        S: VerifierTraceGen<PB, RootSC, DC>,
33        T: RootTraceGen<PB, DC>,
34    {
35        assert_eq!(
36            user_pvs_proof.public_values.len(),
37            self.circuit.num_user_pvs
38        );
39
40        // These AIRs should have the same height regardless of proof or user_pvs_proof.
41        let mut pre_data = self.agg_node_tracegen.generate_pre_verifier_subcircuit_ctx(
42            &proof,
43            user_pvs_proof,
44            self.circuit.memory_dimensions,
45            device_ctx,
46        );
47        let (post_verifier_subcircuit_ctxs, other_compress_inputs) =
48            self.agg_node_tracegen.generate_other_proving_ctxs(
49                &proof,
50                self.circuit.memory_dimensions,
51                deferral_merkle_proofs,
52                device_ctx,
53            );
54        pre_data
55            .poseidon2_compress_inputs
56            .extend(other_compress_inputs);
57
58        // Get the verifier sub-circuit trace heights. If deferrals are enabled, there is
59        // an additional AIR at the end.
60        let verifier_trace_heights = self.trace_heights.as_ref().map(|v| {
61            let num_airs = v.len() - deferral_merkle_proofs.is_some() as usize;
62            &v[3..num_airs]
63        });
64
65        let power_check_inputs = vec![];
66        let mut external_data = VerifierExternalData {
67            poseidon2_compress_inputs: &pre_data.poseidon2_compress_inputs,
68            poseidon2_permute_inputs: &pre_data.poseidon2_permute_inputs,
69            range_check_inputs: &pre_data.range_check_inputs,
70            power_check_inputs: &power_check_inputs,
71            required_heights: verifier_trace_heights,
72            final_transcript_state: None,
73        };
74
75        let subcircuit_ctxs = self.circuit.verifier_circuit.generate_proving_ctxs(
76            &self.child_vk,
77            CachedTraceCtx::Records(self.cached_trace_record.clone()),
78            &[proof],
79            &mut external_data,
80            device_ctx,
81            default_duplex_sponge_recorder(),
82        );
83
84        subcircuit_ctxs.map(|subcircuit_ctxs| ProvingContext {
85            per_trace: pre_data
86                .air_proving_ctxs
87                .into_iter()
88                .chain(subcircuit_ctxs)
89                .chain(post_verifier_subcircuit_ctxs)
90                .enumerate()
91                .collect_vec(),
92        })
93    }
94
95    #[instrument(name = "trace_gen", skip_all)]
96    pub fn generate_proving_ctx_no_def<PB, DC: Clone + Send + Sync>(
97        &self,
98        proof: Proof<SC>,
99        user_pvs_proof: &UserPublicValuesProof<DIGEST_SIZE, PB::Val>,
100        device_ctx: &DC,
101    ) -> Option<ProvingContext<PB>>
102    where
103        PB: ProverBackend<Val = F, Challenge = EF>,
104        PB::Matrix: Clone,
105        S: VerifierTraceGen<PB, RootSC, DC>,
106        T: RootTraceGen<PB, DC>,
107    {
108        assert!(
109            self.circuit.def_hook_commit.is_none(),
110            "deferral-enabled root prover requires generate_proving_ctx_with_deferrals"
111        );
112        self.generate_proving_ctx(proof, user_pvs_proof, None, device_ctx)
113    }
114}