openvm_stark_backend/prover/
types.rs

1use std::{cmp::Reverse, sync::Arc};
2
3use derivative::Derivative;
4
5use crate::{
6    keygen::types::{
7        LinearConstraint, MultiStarkVerifyingKey, MultiStarkVerifyingKey0, StarkVerifyingKey,
8    },
9    proof::TraceVData,
10    prover::{MatrixDimensions, ProverBackend},
11    StarkProtocolConfig, SystemParams,
12};
13
14/// The committed trace data for a single trace matrix. This type is used to store prover data for
15/// both preprocessed trace and cached trace.
16#[derive(Derivative)]
17#[derivative(Clone(bound = "PB::Matrix: Clone"))]
18pub struct CommittedTraceData<PB: ProverBackend> {
19    /// The polynomial commitment.
20    pub commitment: PB::Commitment,
21    /// The trace matrix, unstacked, in evaluation form.
22    pub trace: PB::Matrix,
23    /// The PCS data for a single committed trace matrix.
24    pub data: Arc<PB::PcsData>,
25}
26
27/// The proving key for a circuit consisting of multiple AIRs, after prover-specific data has been
28/// transferred to device. The host data (e.g., vkey) is owned by this struct.
29///
30/// Ordering is always by AIR ID and includes all AIRs, including ones that may have empty traces.
31#[derive(derive_new::new)]
32pub struct DeviceMultiStarkProvingKey<PB: ProverBackend> {
33    pub per_air: Vec<DeviceStarkProvingKey<PB>>,
34    pub trace_height_constraints: Vec<LinearConstraint>,
35    /// Maximum degree of constraints across all AIRs
36    pub max_constraint_degree: usize,
37    pub params: SystemParams,
38    pub vk_pre_hash: PB::Commitment,
39}
40
41/// The proving key after prover-specific data has been transferred to device. The host data (e.g.,
42/// vkey) is owned by this struct.
43pub struct DeviceStarkProvingKey<PB: ProverBackend> {
44    /// Type name of the AIR, for display purposes only
45    pub air_name: String,
46    pub vk: StarkVerifyingKey<PB::Val, PB::Commitment>,
47    /// Prover only data for preprocessed trace
48    pub preprocessed_data: Option<CommittedTraceData<PB>>,
49    pub other_data: PB::OtherAirData,
50}
51
52#[derive(derive_new::new)]
53pub struct ProvingContext<PB: ProverBackend> {
54    /// For each AIR with non-empty trace, the pair of (AIR ID, [AirProvingContext]), where AIR
55    /// ID is with respect to the vkey ordering.
56    pub per_trace: Vec<(usize, AirProvingContext<PB>)>,
57}
58
59#[derive(derive_new::new)]
60pub struct AirProvingContext<PB: ProverBackend> {
61    /// Cached main trace matrices as `PcsData`. The original trace matrix should be extractable as
62    /// a view from the `PcsData`. The `PcsData` should also contain the commitment value. Cached
63    /// trace commitments have a single matrix per commitment.
64    ///
65    /// The `PcsData` is kept inside an `Arc` to emphasize that this data is cached and may be
66    /// shared between multiple proving contexts. In particular, it is not typically safe to mutate
67    /// the data during a proving job.
68    pub cached_mains: Vec<CommittedTraceData<PB>>,
69    /// Common main trace matrix
70    pub common_main: PB::Matrix,
71    /// Public values
72    pub public_values: Vec<PB::Val>,
73}
74
75/// Proof on the host, with respect to the host types in the generic `PB`.
76pub struct HostProof<SC: StarkProtocolConfig, PB: ProverBackend, ConstraintsProof, OpeningProof> {
77    /// The commitment to the data in common_main.
78    pub common_main_commit: PB::Commitment,
79
80    /// For each AIR in vkey order, the corresponding trace shape, or None if
81    /// the trace is empty. In a valid proof, if `vk.per_air[i].is_required`,
82    /// then `trace_vdata[i]` must be `Some(_)`.
83    pub trace_vdata: Vec<Option<TraceVData<SC>>>,
84
85    /// For each AIR in vkey order, the public values. Public values should be empty if the AIR has
86    /// an empty trace.
87    pub public_values: Vec<Vec<PB::Val>>,
88
89    pub constraints_proof: ConstraintsProof,
90    /// Opening proof for multiple polynomials over mixed sized domains
91    pub opening_proof: OpeningProof,
92}
93
94impl<PB: ProverBackend> CommittedTraceData<PB> {
95    #[inline(always)]
96    pub fn height(&self) -> usize {
97        self.trace.height()
98    }
99}
100
101impl<PB: ProverBackend> DeviceMultiStarkProvingKey<PB> {
102    pub fn get_vk<SC>(&self) -> MultiStarkVerifyingKey<SC>
103    where
104        SC: StarkProtocolConfig<F = PB::Val, Digest = PB::Commitment>,
105    {
106        let per_air = self.per_air.iter().map(|pk| pk.vk.clone()).collect();
107        let inner = MultiStarkVerifyingKey0 {
108            params: self.params.clone(),
109            per_air,
110            trace_height_constraints: self.trace_height_constraints.clone(),
111        };
112        MultiStarkVerifyingKey {
113            inner,
114            pre_hash: self.vk_pre_hash.clone(),
115        }
116    }
117}
118
119impl<PB: ProverBackend> IntoIterator for ProvingContext<PB> {
120    type Item = (usize, AirProvingContext<PB>);
121    type IntoIter = std::vec::IntoIter<Self::Item>;
122
123    fn into_iter(self) -> Self::IntoIter {
124        self.per_trace.into_iter()
125    }
126}
127
128impl<PB: ProverBackend> ProvingContext<PB> {
129    pub fn common_main_traces(&self) -> impl Iterator<Item = (usize, &PB::Matrix)> {
130        self.per_trace
131            .iter()
132            .map(|(air_idx, trace_ctx)| (*air_idx, &trace_ctx.common_main))
133    }
134
135    // Returns `self` with the trace data sorted to be descending in height for column stacking. For
136    // equal heights, traces are sorted in ascending order of AIR index.
137    pub fn into_sorted(mut self) -> Self {
138        self.sort_for_stacking();
139        self
140    }
141
142    // Stable sort the trace data to be descending in height: this is needed for stacking. For
143    // equal heights, sort in ascending order of AIR index.
144    pub fn sort_for_stacking(&mut self) {
145        self.per_trace.sort_by_key(|(air_idx, trace_ctx)| {
146            (Reverse(trace_ctx.common_main.height()), *air_idx)
147        });
148    }
149}
150
151impl<PB: ProverBackend> AirProvingContext<PB> {
152    pub fn simple(common_main_trace: PB::Matrix, public_values: Vec<PB::Val>) -> Self {
153        Self::new(vec![], common_main_trace, public_values)
154    }
155    pub fn simple_no_pis(common_main_trace: PB::Matrix) -> Self {
156        Self::simple(common_main_trace, vec![])
157    }
158
159    /// Return the height of the main trace.
160    pub fn height(&self) -> usize {
161        self.common_main.height()
162    }
163}
openvm_stark_backend/prover/types.rs

openvm_stark_backend/prover/
types.rs
