openvm_stark_backend/prover/cpu/quotient/
mod.rs

1use std::sync::Arc;
2
3use itertools::{izip, multiunzip, Itertools};
4use p3_commit::{Pcs, PolynomialSpace};
5use p3_field::FieldAlgebra;
6use p3_matrix::{dense::RowMajorMatrix, Matrix};
7use p3_util::log2_strict_usize;
8use tracing::instrument;
9
10use self::single::compute_single_rap_quotient_values;
11use super::PcsData;
12use crate::{
13    air_builders::symbolic::SymbolicExpressionDag,
14    config::{Com, Domain, PackedChallenge, StarkGenericConfig, Val},
15    prover::types::RapView,
16};
17
18mod evaluator;
19pub(crate) mod single;
20
21pub struct QuotientCommitter<'pcs, SC: StarkGenericConfig> {
22    pcs: &'pcs SC::Pcs,
23    alpha: SC::Challenge,
24    extra_capacity_bits: usize,
25}
26
27impl<'pcs, SC: StarkGenericConfig> QuotientCommitter<'pcs, SC> {
28    pub fn new(pcs: &'pcs SC::Pcs, alpha: SC::Challenge, extra_capacity_bits: usize) -> Self {
29        Self {
30            pcs,
31            alpha,
32            extra_capacity_bits,
33        }
34    }
35
36    /// Constructs quotient domains and computes the evaluation of the quotient polynomials
37    /// on the quotient domains of each RAP.
38    ///
39    /// ## Assumptions
40    /// - `constraints`, `extended_views`, `quotient_degrees` have equal lengths and the length
41    ///   equals number of RAPs.
42    /// - `quotient_degrees` is the factor to **multiply** the trace degree by to get the degree of
43    ///   the quotient polynomial. This should be determined from the constraint degree of the RAP.
44    /// - `extended_views` is a view of the trace polynomials evaluated on the quotient domain, with
45    ///   rows bit reversed to account for the fact that the quotient domain is different for each
46    ///   RAP.
47    ///
48    /// **Note**: This function assumes that the
49    /// `quotient_domain.split_evals(quotient_degree, quotient_flat)` function from Plonky3 works
50    /// as described in [compute_single_rap_quotient_values].
51    #[instrument(name = "quotient_poly_compute", level = "info", skip_all)]
52    pub fn quotient_values(
53        &self,
54        constraints: &[&SymbolicExpressionDag<Val<SC>>],
55        extended_views: Vec<RapView<impl Matrix<Val<SC>>, Val<SC>, SC::Challenge>>,
56        quotient_degrees: &[u8],
57    ) -> QuotientData<SC> {
58        let max_alpha_pow = constraints
59            .iter()
60            .map(|c| c.constraint_idx.len())
61            .max()
62            .unwrap_or(0);
63        let alpha_powers = self
64            .alpha
65            .powers()
66            .take(max_alpha_pow)
67            .map(PackedChallenge::<SC>::from_f)
68            .collect_vec();
69        assert_eq!(constraints.len(), extended_views.len());
70        assert_eq!(constraints.len(), quotient_degrees.len());
71        let chunks = izip!(constraints, extended_views, quotient_degrees)
72            .flat_map(|(constraints, extended_view, &quotient_degree)| {
73                self.single_rap_quotient_values(
74                    constraints,
75                    extended_view,
76                    quotient_degree,
77                    &alpha_powers,
78                )
79            })
80            .collect();
81        QuotientData { chunks }
82    }
83
84    pub(super) fn single_rap_quotient_values(
85        &self,
86        constraints: &SymbolicExpressionDag<Val<SC>>,
87        view: RapView<impl Matrix<Val<SC>>, Val<SC>, SC::Challenge>,
88        quotient_degree: u8,
89        alpha_powers: &[PackedChallenge<SC>],
90    ) -> impl IntoIterator<Item = QuotientChunk<SC>> {
91        let log_trace_height = view.log_trace_height;
92        let trace_domain = self
93            .pcs
94            .natural_domain_for_degree(1usize << log_trace_height);
95        let quotient_domain =
96            trace_domain.create_disjoint_domain(trace_domain.size() * quotient_degree as usize);
97
98        let (after_challenge_lde_on_quotient_domain, challenges, exposed_values_after_challenge): (
99            Vec<_>,
100            Vec<_>,
101            Vec<_>,
102        ) = multiunzip(view.per_phase.into_iter().map(|view| {
103            (
104                view.inner
105                    .expect("gap in challenge phase not supported yet"),
106                view.challenges
107                    .into_iter()
108                    .map(PackedChallenge::<SC>::from_f)
109                    .collect_vec(),
110                view.exposed_values
111                    .into_iter()
112                    .map(PackedChallenge::<SC>::from_f)
113                    .collect_vec(),
114            )
115        }));
116
117        compute_single_rap_quotient_values::<SC, _>(
118            constraints,
119            trace_domain,
120            quotient_domain,
121            view.preprocessed,
122            view.partitioned_main,
123            after_challenge_lde_on_quotient_domain,
124            &challenges,
125            alpha_powers,
126            &view.public_values,
127            &exposed_values_after_challenge,
128            self.extra_capacity_bits,
129        )
130    }
131
132    #[instrument(name = "quotient_poly_commit", skip_all)]
133    pub fn commit(&self, data: QuotientData<SC>) -> (Com<SC>, PcsData<SC>) {
134        let (log_trace_heights, quotient_domains_and_chunks): (Vec<_>, Vec<_>) = data
135            .chunks
136            .into_iter()
137            .map(|q| {
138                (
139                    log2_strict_usize(q.domain.size()) as u8,
140                    (q.domain, q.matrix),
141                )
142            })
143            .unzip();
144        let (commit, data) = self.pcs.commit(quotient_domains_and_chunks);
145        (
146            commit,
147            PcsData {
148                data: Arc::new(data),
149                log_trace_heights,
150            },
151        )
152    }
153}
154
155/// The quotient polynomials from multiple RAP matrices.
156pub struct QuotientData<SC: StarkGenericConfig> {
157    /// Length equals `quotient_ degree`.
158    chunks: Vec<QuotientChunk<SC>>,
159}
160
161/// The vector of evaluations of the quotient polynomial on the quotient domain,
162/// split into chunks of size equal to the trace domain size (quotient domain size
163/// divided by `quotient_degree`).
164///
165/// This represents a single chunk, where the vector of extension field elements is
166/// further flattened to a matrix of base field elements.
167pub struct QuotientChunk<SC: StarkGenericConfig> {
168    /// Chunk of quotient domain, which is a coset of the trace domain
169    pub domain: Domain<SC>,
170    /// Matrix with number of rows equal to trace domain size,
171    /// and number of columns equal to extension field degree.
172    pub matrix: RowMajorMatrix<Val<SC>>,
173}