openvm_stark_backend/prover/
mod.rs

1use itertools::{izip, Itertools};
2use p3_field::PrimeCharacteristicRing;
3use p3_util::log2_strict_usize;
4use tracing::{info, info_span, instrument};
5
6#[cfg(feature = "metrics")]
7use crate::prover::metrics::trace_metrics;
8use crate::{
9    proof::{BatchConstraintProof, GkrProof, Proof, StackingProof, TraceVData, WhirProof},
10    FiatShamirTranscript, StarkProtocolConfig,
11};
12
13mod cpu_backend;
14pub mod error;
15mod hal;
16mod logup_zerocheck;
17mod matrix;
18pub mod metrics;
19pub mod poly;
20pub mod stacked_pcs;
21pub mod stacked_reduction;
22pub mod sumcheck;
23mod types;
24pub mod whir;
25
26pub use cpu_backend::*;
27pub use error::*;
28pub use hal::*;
29pub use logup_zerocheck::*;
30pub use matrix::*;
31pub use types::*;
32
33/// Trait for STARK/SNARK proving at the highest abstraction level.
34pub trait Prover {
35    type ProvingKeyView<'a>
36    where
37        Self: 'a;
38    type ProvingContext<'a>
39    where
40        Self: 'a;
41    type Proof;
42    type Error;
43
44    /// The prover should own the challenger, whose state mutates during proving.
45    fn prove<'a>(
46        &'a mut self,
47        pk: Self::ProvingKeyView<'a>,
48        ctx: Self::ProvingContext<'a>,
49    ) -> Result<Self::Proof, Self::Error>;
50}
51
52pub struct Coordinator<SC: StarkProtocolConfig, PB: ProverBackend, PD, TS> {
53    pub backend: PB,
54    pub device: PD,
55    pub(crate) transcript: TS,
56    _sc: std::marker::PhantomData<SC>,
57}
58
59impl<SC: StarkProtocolConfig, PB: ProverBackend, PD, TS> Coordinator<SC, PB, PD, TS> {
60    pub fn new(backend: PB, device: PD, transcript: TS) -> Self {
61        Self {
62            backend,
63            device,
64            transcript,
65            _sc: std::marker::PhantomData,
66        }
67    }
68}
69
70impl<SC, PB, PD, TS> Prover for Coordinator<SC, PB, PD, TS>
71where
72    SC: StarkProtocolConfig,
73    PB: ProverBackend<Val = SC::F, Challenge = SC::EF, Commitment = SC::Digest>,
74    PD: ProverDevice<PB, TS>,
75    PD::Artifacts: Into<PD::OpeningPoints>,
76    PD::PartialProof: Into<(GkrProof<SC>, BatchConstraintProof<SC>)>,
77    PD::OpeningProof: Into<(StackingProof<SC>, WhirProof<SC>)>,
78    TS: FiatShamirTranscript<SC>,
79{
80    type Proof = Proof<SC>;
81    type Error = <PD as ProverDevice<PB, TS>>::Error;
82    type ProvingKeyView<'a>
83        = &'a DeviceMultiStarkProvingKey<PB>
84    where
85        Self: 'a;
86
87    type ProvingContext<'a>
88        = ProvingContext<PB>
89    where
90        Self: 'a;
91
92    /// Specialized prove for InteractiveAirs.
93    /// Handles trace generation of the permutation traces.
94    /// Assumes the main traces have been generated and committed already.
95    ///
96    /// The [DeviceMultiStarkProvingKey] should already be filtered to only include the relevant
97    /// AIR's proving keys.
98    #[instrument(
99        name = "stark_prove_excluding_trace",
100        level = "info",
101        skip_all,
102        fields(phase = "prover")
103    )]
104    fn prove<'a>(
105        &'a mut self,
106        mpk: &'a DeviceMultiStarkProvingKey<PB>,
107        unsorted_ctx: ProvingContext<PB>,
108    ) -> Result<Self::Proof, Self::Error> {
109        let transcript = &mut self.transcript;
110        transcript.observe_commit(mpk.vk_pre_hash);
111
112        let ctx = unsorted_ctx.into_sorted();
113        // `ctx` should NOT be permuted anymore: the ordering by `trace_idx` is now fixed.
114
115        let num_airs_present = ctx.per_trace.len();
116        info!(num_airs_present);
117
118        let _main_commit_span = info_span!("prover.main_trace_commit", phase = "prover").entered();
119        let (common_main_commit, common_main_pcs_data) = {
120            let traces = ctx
121                .common_main_traces()
122                .map(|(_, trace)| trace)
123                .collect_vec();
124            self.device.commit(&traces)?
125        };
126
127        let mut trace_vdata: Vec<Option<TraceVData<SC>>> = vec![None; mpk.per_air.len()];
128        let mut public_values: Vec<Vec<SC::F>> = vec![Vec::new(); mpk.per_air.len()];
129
130        // Hypercube dimension per trace (present AIR)
131        for (air_id, trace_ctx) in &ctx.per_trace {
132            let trace_height = trace_ctx.common_main.height();
133            let log_height = log2_strict_usize(trace_height);
134
135            trace_vdata[*air_id] = Some(TraceVData::<SC> {
136                log_height,
137                cached_commitments: trace_ctx
138                    .cached_mains
139                    .iter()
140                    .map(|cd| cd.commitment)
141                    .collect(),
142            });
143            public_values[*air_id] = trace_ctx.public_values.clone();
144        }
145        #[cfg(feature = "metrics")]
146        trace_metrics::<SC, _>(mpk, &trace_vdata).emit();
147
148        // Only observe commits for present AIRs.
149        // Commitments order:
150        // - 1 commitment of all common main traces
151        // - for each air:
152        //   - preprocessed commit if present
153        //   - for each cached main trace
154        //     - 1 commitment
155        transcript.observe_commit(common_main_commit);
156        drop(_main_commit_span);
157
158        for (trace_vdata, pvs, pk) in izip!(&trace_vdata, &public_values, &mpk.per_air) {
159            if !pk.vk.is_required {
160                transcript.observe(SC::F::from_bool(trace_vdata.is_some()));
161            }
162            if let Some(trace_vdata) = trace_vdata {
163                if let Some(cd) = &pk.preprocessed_data {
164                    transcript.observe_commit(cd.commitment);
165                } else {
166                    transcript.observe(SC::F::from_usize(trace_vdata.log_height));
167                }
168                for commit in &trace_vdata.cached_commitments {
169                    transcript.observe_commit(*commit);
170                }
171            }
172            for pv in pvs {
173                transcript.observe(*pv);
174            }
175        }
176
177        let (constraints_proof, r) =
178            self.device
179                .prove_rap_constraints(transcript, mpk, &ctx, &common_main_pcs_data)?;
180
181        let opening_proof =
182            self.device
183                .prove_openings(transcript, mpk, ctx, common_main_pcs_data, r.into())?;
184
185        let (gkr_proof, batch_constraint_proof) = constraints_proof.into();
186        let (stacking_proof, whir_proof) = opening_proof.into();
187
188        Ok(Proof::<SC> {
189            public_values,
190            trace_vdata,
191            common_main_commit,
192            gkr_proof,
193            batch_constraint_proof,
194            stacking_proof,
195            whir_proof,
196        })
197    }
198}
openvm_stark_backend/prover/mod.rs

openvm_stark_backend/prover/
mod.rs
