openvm_stark_backend/prover/
types.rs

1use std::sync::Arc;
2
3use derivative::Derivative;
4use p3_field::Field;
5use p3_matrix::{dense::RowMajorMatrix, Matrix};
6use serde::{Deserialize, Serialize};
7
8use super::hal::ProverBackend;
9use crate::{
10    config::{Com, PcsProof, PcsProverData, RapPhaseSeqPartialProof, StarkGenericConfig, Val},
11    keygen::types::{LinearConstraint, StarkVerifyingKey},
12    proof::{AirProofData, Commitments, OpeningProof, Proof},
13};
14
15/// A view of the proving key after it has been transferred to device.
16pub struct DeviceMultiStarkProvingKey<'a, PB: ProverBackend> {
17    pub(super) air_ids: Vec<usize>,
18    pub per_air: Vec<DeviceStarkProvingKey<'a, PB>>,
19    /// Each [LinearConstraint] is indexed by AIR ID.
20    /// **Caution**: the linear constraints are **not** filtered for only the AIRs appearing in `per_air`.
21    pub trace_height_constraints: Vec<LinearConstraint>,
22    pub vk_pre_hash: PB::Commitment,
23}
24
25impl<'a, PB: ProverBackend> DeviceMultiStarkProvingKey<'a, PB> {
26    pub fn new(
27        air_ids: Vec<usize>,
28        per_air: Vec<DeviceStarkProvingKey<'a, PB>>,
29        trace_height_constraints: Vec<LinearConstraint>,
30        vk_pre_hash: PB::Commitment,
31    ) -> Self {
32        assert_eq!(air_ids.len(), per_air.len());
33        Self {
34            air_ids,
35            per_air,
36            trace_height_constraints,
37            vk_pre_hash,
38        }
39    }
40}
41
42pub struct DeviceStarkProvingKey<'a, PB: ProverBackend> {
43    /// Type name of the AIR, for display purposes only
44    pub air_name: &'a str,
45    pub vk: &'a StarkVerifyingKey<PB::Val, PB::Commitment>,
46    /// Prover only data for preprocessed trace
47    pub preprocessed_data: Option<SingleCommitPreimage<PB::Matrix, PB::PcsData>>,
48    /// Additional configuration or preprocessed data for the RAP phases
49    pub rap_partial_pk: PB::RapPartialProvingKey,
50}
51
52/// A view of an already committed trace, together with a reference to the
53/// preimage of the commitment.
54///
55/// The PCS may commit to multiple matrices at once, so `matrix_idx` provides
56/// the index of this matrix within the commitment.
57#[derive(Clone)]
58pub struct SingleCommitPreimage<Matrix, PcsData> {
59    pub trace: Matrix,
60    pub data: PcsData,
61    pub matrix_idx: u32,
62}
63
64#[derive(derive_new::new)]
65pub struct ProvingContext<'a, PB: ProverBackend> {
66    /// (AIR id, AIR input)
67    pub per_air: Vec<(usize, AirProvingContext<'a, PB>)>,
68}
69
70impl<'a, PB: ProverBackend> ProvingContext<'a, PB> {
71    pub fn into_air_proving_ctx_vec(self) -> Vec<AirProvingContext<'a, PB>> {
72        self.per_air.into_iter().map(|(_, x)| x).collect()
73    }
74}
75
76impl<'a, PB: ProverBackend> IntoIterator for ProvingContext<'a, PB> {
77    type Item = (usize, AirProvingContext<'a, PB>);
78    type IntoIter = std::vec::IntoIter<Self::Item>;
79
80    fn into_iter(self) -> Self::IntoIter {
81        self.per_air.into_iter()
82    }
83}
84
85/// Necessary context for proving a single AIR.
86/// Consists of the trace and public values (witness and instance data) as well as
87/// cached prover data from cached traces (already committed traces).
88///
89/// Public values: for each AIR, a separate list of public values.
90/// The prover can support global public values that are shared among all AIRs,
91/// but we currently split public values per-AIR for modularity.
92pub struct AirProvingContext<'a, PB: ProverBackend> {
93    /// Cached main trace matrices with commitments. One matrix per commitment.
94    #[allow(clippy::type_complexity)]
95    pub cached_mains: Vec<(
96        PB::Commitment,
97        SingleCommitPreimage<&'a PB::Matrix, &'a PB::PcsData>,
98    )>,
99    /// Common main trace matrix
100    pub common_main: Option<PB::Matrix>,
101    /// Public values
102    // [jpw] This is on host for now because it seems more convenient for the challenger to be on host.
103    pub public_values: Vec<PB::Val>,
104}
105
106/// A view of just the preprocessed AIR (PAIR), without any after challenge columns.
107/// The PAIR trace consists of horizontal concatenation of multiple matrices of the same height:
108/// - preprocessed trace matrix
109/// - the main trace matrix is horizontally partitioned into multiple matrices,
110///   where each matrix can belong to a separate matrix commitment.
111///
112/// The generic `T` may be either just the trace matrix view or the LDE matrix view.
113pub struct PairView<T, Val> {
114    /// Log_2 of the trace domain size (i.e., height of matrices)
115    pub log_trace_height: u8,
116    /// Preprocessed trace data, if any
117    pub preprocessed: Option<T>,
118    /// Main trace data, horizontally partitioned into multiple matrices
119    pub partitioned_main: Vec<T>,
120    /// Public values
121    pub public_values: Vec<Val>,
122}
123
124/// The full RAP trace consists of horizontal concatenation of multiple matrices of the same height:
125/// - preprocessed trace matrix
126/// - the main trace matrix is horizontally partitioned into multiple matrices,
127///   where each matrix can belong to a separate matrix commitment.
128/// - after each round of challenges, a trace matrix for trace allowed to use those challenges
129///
130/// Each of these matrices is allowed to be in a separate commitment.
131///
132/// Only the main trace matrix is allowed to be partitioned, so that different parts may belong to
133/// different commitments. We do not see any use cases where the `preprocessed` or `after_challenge`
134/// matrices need to be partitioned.
135///
136/// The generic `T` may be either just the trace matrix view or the LDE matrix view.
137pub struct RapView<T, Val, Challenge> {
138    pub pair: PairView<T, Val>,
139    /// `per_phase[i]` is a view which is calculated after sampling challenges
140    /// which depend on observing commitments to `pair` and `per_phase[..i]`.
141    pub per_phase: Vec<RapSinglePhaseView<T, Challenge>>,
142}
143
144#[derive(Clone)]
145pub struct RapSinglePhaseView<T, Challenge> {
146    /// `None` if this RAP does not use this phase.
147    pub inner: Option<T>,
148    /// The challenges sampled in this phase
149    pub challenges: Vec<Challenge>,
150    /// These are extension field values exposed to the verifier after this phase.
151    /// They are like public values, except that they depend on randomness.
152    pub exposed_values: Vec<Challenge>,
153}
154
155impl<T, Challenge> Default for RapSinglePhaseView<T, Challenge> {
156    fn default() -> Self {
157        Self {
158            inner: None,
159            challenges: Vec::new(),
160            exposed_values: Vec::new(),
161        }
162    }
163}
164
165#[derive(derive_new::new)]
166pub struct ProverDataAfterRapPhases<PB: ProverBackend> {
167    /// For each challenge phase **after** the main phase,
168    /// the commitment and preimage (there should never be a reason to have more than one).
169    /// This may be empty if challenge phases do not require additional trace commitments.
170    pub committed_pcs_data_per_phase: Vec<(PB::Commitment, PB::PcsData)>,
171    /// For each challenge phase, for each RAP,
172    /// the challenge, and exposed values for the RAP.
173    /// The indexing is `rap_views_per_phase[phase_idx][rap_idx]`.
174    pub rap_views_per_phase: Vec<Vec<RapSinglePhaseView<usize, PB::Challenge>>>,
175}
176
177/// The full proof for multiple RAPs where trace matrices are committed into
178/// multiple commitments, where each commitment is multi-matrix.
179///
180/// Includes the quotient commitments and FRI opening proofs for the constraints as well.
181#[derive(Serialize, Deserialize, Derivative)]
182#[serde(bound = "")]
183#[derivative(Clone(bound = ""))]
184pub struct HalProof<PB: ProverBackend> {
185    /// The PCS commitments
186    pub commitments: Commitments<PB::Commitment>,
187    /// Opening proofs separated by partition, but this may change
188    pub opening: PB::OpeningProof,
189    /// Proof data for each AIR
190    pub per_air: Vec<AirProofData<PB::Val, PB::Challenge>>,
191    /// Partial proof for rap phase if it exists
192    pub rap_partial_proof: PB::RapPartialProof,
193}
194
195impl<PB, SC: StarkGenericConfig> From<HalProof<PB>> for Proof<SC>
196where
197    PB: ProverBackend<Val = Val<SC>, Challenge = SC::Challenge, Commitment = Com<SC>>,
198    PB::OpeningProof: Into<OpeningProof<PcsProof<SC>, SC::Challenge>>,
199    PB::RapPartialProof: Into<Option<RapPhaseSeqPartialProof<SC>>>,
200{
201    fn from(proof: HalProof<PB>) -> Self {
202        Proof {
203            commitments: proof.commitments,
204            opening: proof.opening.into(),
205            per_air: proof.per_air,
206            rap_phase_seq_proof: proof.rap_partial_proof.into(),
207        }
208    }
209}
210
211// ============= Below are common types independent of hardware ============
212
213#[derive(Derivative, derive_new::new)]
214#[derivative(Clone(bound = "Com<SC>: Clone"))]
215pub struct ProofInput<SC: StarkGenericConfig> {
216    /// (AIR id, AIR input)
217    pub per_air: Vec<(usize, AirProofInput<SC>)>,
218}
219
220#[derive(Serialize, Deserialize, Derivative)]
221#[serde(bound(
222    serialize = "PcsProverData<SC>: Serialize",
223    deserialize = "PcsProverData<SC>: Deserialize<'de>"
224))]
225#[derivative(Clone(bound = "Com<SC>: Clone"))]
226pub struct CommittedTraceData<SC: StarkGenericConfig> {
227    pub trace: Arc<RowMajorMatrix<Val<SC>>>,
228    pub commitment: Com<SC>,
229    pub pcs_data: Arc<PcsProverData<SC>>,
230}
231
232/// Necessary input for proving a single AIR.
233///
234/// The [Chip](crate::chip::Chip) trait is currently specific to the
235/// CPU backend and in particular to `RowMajorMatrix`. We may extend
236/// to more general [ProverBackend](super::hal::ProverBackend)s, but
237/// currently we use this struct as a common interface.
238#[derive(Derivative)]
239#[derivative(Clone(bound = "Com<SC>: Clone"))]
240pub struct AirProofInput<SC: StarkGenericConfig> {
241    /// Prover data for cached main traces.
242    /// They must either be all provided or they will be regenerated
243    /// from the raw traces.
244    pub cached_mains_pdata: Vec<(Com<SC>, Arc<PcsProverData<SC>>)>,
245    pub raw: AirProofRawInput<Val<SC>>,
246}
247
248/// Raw input for proving a single AIR.
249#[derive(Clone, Debug)]
250pub struct AirProofRawInput<F: Field> {
251    /// Cached main trace matrices
252    pub cached_mains: Vec<Arc<RowMajorMatrix<F>>>,
253    /// Common main trace matrix
254    pub common_main: Option<RowMajorMatrix<F>>,
255    /// Public values
256    pub public_values: Vec<F>,
257}
258
259impl<F: Field> AirProofRawInput<F> {
260    pub fn height(&self) -> usize {
261        let mut height = None;
262        for m in self.cached_mains.iter() {
263            if let Some(h) = height {
264                assert_eq!(h, m.height());
265            } else {
266                height = Some(m.height());
267            }
268        }
269        let common_h = self.common_main.as_ref().map(|trace| trace.height());
270        if let Some(h) = height {
271            if let Some(common_h) = common_h {
272                assert_eq!(h, common_h);
273            }
274            h
275        } else {
276            common_h.unwrap_or(0)
277        }
278    }
279}
openvm_stark_backend/prover/types.rs

openvm_stark_backend/prover/
types.rs
