openvm_stark_backend/prover/
stacked_pcs.rs

1use getset::{CopyGetters, Getters};
2use itertools::Itertools;
3use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft};
4use p3_field::{ExtensionField, Field, TwoAdicField};
5use p3_maybe_rayon::prelude::*;
6use p3_util::log2_strict_usize;
7use serde::{Deserialize, Serialize};
8use tracing::instrument;
9
10use crate::{
11    hasher::MerkleHasher,
12    prover::{
13        col_maj_idx, error::StackedPcsError, poly::eval_to_coeff_rs_message, ColMajorMatrix,
14        MatrixDimensions, MatrixView, StridedColMajorMatrixView,
15    },
16};
17
18#[derive(Clone, Serialize, Deserialize, Debug, CopyGetters)]
19pub struct StackedLayout {
20    /// The minimum log2 height of a stacked slice. When stacking columns with smaller height, the
21    /// column is expanded to `2^l_skip` by striding.
22    #[getset(get_copy = "pub")]
23    l_skip: usize,
24    /// Stacked height
25    #[getset(get_copy = "pub")]
26    height: usize,
27    /// Stacked width
28    #[getset(get_copy = "pub")]
29    width: usize,
30    /// The columns of the unstacked matrices in sorted order. Each entry `(matrix index, column
31    /// index, coordinate)` contains the pointer `(matrix index, column index)` to a column of the
32    /// unstacked collection of matrices as well as `coordinate` which is a pointer to where the
33    /// column starts in the stacked matrix.
34    pub sorted_cols: Vec<(
35        usize, /* unstacked matrix index */
36        usize, /* unstacked column index */
37        StackedSlice,
38    )>,
39    /// `mat_starts[mat_idx]` is the index in `sorted_cols` where the matrix with index `mat_idx`
40    /// starts.
41    pub mat_starts: Vec<usize>,
42}
43
44/// Pointer to the location of a sub-column within the stacked matrix.
45/// This struct contains length information, but information from [StackedLayout] (namely `l_skip`)
46/// is needed to determine if this is a strided slice or not.
47#[derive(Copy, Clone, Debug, Serialize, Deserialize, CopyGetters, derive_new::new)]
48pub struct StackedSlice {
49    pub col_idx: usize,
50    pub row_idx: usize,
51    /// The true log height. If `>= l_skip`, no striding. Otherwise striding by `2^{l_skip -
52    /// log_height}`.
53    #[getset(get_copy = "pub")]
54    log_height: usize,
55}
56
57impl StackedSlice {
58    #[inline(always)]
59    pub fn len(&self, l_skip: usize) -> usize {
60        Self::_len(self.log_height, l_skip)
61    }
62
63    #[inline(always)]
64    pub fn stride(&self, l_skip: usize) -> usize {
65        1 << l_skip.saturating_sub(self.log_height)
66    }
67
68    #[inline(always)]
69    fn _len(log_height: usize, l_skip: usize) -> usize {
70        if l_skip <= log_height {
71            1 << log_height
72        } else {
73            1 << l_skip
74        }
75    }
76}
77
78#[derive(Clone, Debug, Getters, CopyGetters, Serialize, Deserialize)]
79pub struct MerkleTree<F, Digest> {
80    /// The matrix that is used to form the leaves of the Merkle tree, which are
81    /// in turn hashed into the bottom digest layer.
82    ///
83    /// This is typically the codeword matrix in hash-based PCS.
84    #[getset(get = "pub")]
85    pub(crate) backing_matrix: ColMajorMatrix<F>,
86    #[getset(get = "pub")]
87    pub(crate) digest_layers: Vec<Vec<Digest>>,
88    #[getset(get_copy = "pub")]
89    pub(crate) rows_per_query: usize,
90}
91
92#[derive(Clone, Serialize, Deserialize, derive_new::new)]
93pub struct StackedPcsData<F, Digest> {
94    /// Layout of the unstacked collection of matrices within the stacked matrix.
95    pub layout: StackedLayout,
96    /// The stacked matrix of evaluations with height `2^{l_skip + n_stack}`.
97    pub matrix: ColMajorMatrix<F>,
98    /// Merkle tree of the Reed-Solomon codewords of the stacked matrix.
99    /// Depends on `k_whir` parameter.
100    pub tree: MerkleTree<F, Digest>,
101}
102
103impl<F, Digest: Clone> StackedPcsData<F, Digest> {
104    /// Returns the root of the Merkle tree.
105    pub fn commit(&self) -> Result<Digest, StackedPcsError> {
106        self.tree.root()
107    }
108
109    pub fn mat_view(&self, unstacked_mat_idx: usize) -> StridedColMajorMatrixView<'_, F> {
110        self.layout.mat_view(unstacked_mat_idx, &self.matrix)
111    }
112}
113
114#[instrument(level = "info", skip_all)]
115#[allow(clippy::type_complexity)]
116pub fn stacked_commit<H: MerkleHasher>(
117    hasher: &H,
118    l_skip: usize,
119    n_stack: usize,
120    log_blowup: usize,
121    k_whir: usize,
122    traces: &[&ColMajorMatrix<H::F>],
123) -> Result<(H::Digest, StackedPcsData<H::F, H::Digest>), StackedPcsError>
124where
125    H::F: TwoAdicField + Ord,
126    H::Digest: Copy,
127{
128    let (q_trace, layout) = stacked_matrix(l_skip, n_stack, traces)?;
129    let rs_matrix = rs_code_matrix(l_skip, log_blowup, &q_trace)?;
130    let tree = MerkleTree::new(hasher, rs_matrix, 1 << k_whir)?;
131    let root = tree.root()?;
132    let data = StackedPcsData::new(layout, q_trace, tree);
133    Ok((root, data))
134}
135
136impl StackedLayout {
137    /// Computes the layout of greedily stacking columns with dimension metadata given by `sorted`
138    /// into a stacked matrix.
139    /// - `l_skip` is a threshold log2 height: if a column has height less than `2^l_skip`, it is
140    ///   stacked as a column of height `2^l_skip` with stride `2^{l_skip - log_height}`.
141    /// - `log_stacked_height` is the log2 height of the stacked matrix.
142    /// - `sorted` is Vec of `(width, log_height)` that must already be **sorted** in descending
143    ///   order of `log_height`.
144    pub fn new(
145        l_skip: usize,
146        log_stacked_height: usize,
147        sorted: Vec<(usize /* width */, usize /* log_height */)>,
148    ) -> Result<Self, StackedPcsError> {
149        debug_assert!(l_skip <= log_stacked_height);
150        debug_assert!(sorted.is_sorted_by(|a, b| a.1 >= b.1));
151        let mut sorted_cols = Vec::with_capacity(sorted.len());
152        let mut mat_starts = Vec::new();
153        let mut col_idx = 0;
154        let mut row_idx = 0;
155        for (mat_idx, (width, log_ht)) in sorted.into_iter().enumerate() {
156            mat_starts.push(sorted_cols.len());
157            if width == 0 {
158                continue;
159            }
160            if log_ht > log_stacked_height {
161                return Err(StackedPcsError::LayoutHeightExceeded {
162                    log_height: log_ht,
163                    log_stacked_height,
164                });
165            }
166            for j in 0..width {
167                let slice_len = StackedSlice::_len(log_ht, l_skip);
168                if row_idx + slice_len > (1 << log_stacked_height) {
169                    if row_idx != 1 << log_stacked_height {
170                        return Err(StackedPcsError::LayoutRowOverflow {
171                            col_idx,
172                            stacked_height: 1 << log_stacked_height,
173                        });
174                    }
175                    col_idx += 1;
176                    row_idx = 0;
177                }
178                let slice = StackedSlice {
179                    col_idx,
180                    row_idx,
181                    log_height: log_ht,
182                };
183                sorted_cols.push((mat_idx, j, slice));
184                row_idx += slice_len;
185            }
186        }
187        let stacked_width = col_idx + usize::from(row_idx != 0);
188        debug_assert_eq!(
189            stacked_width,
190            sorted_cols
191                .iter()
192                .map(|(_, _, slice)| slice.col_idx + 1)
193                .max()
194                .unwrap_or(0)
195        );
196        Ok(Self {
197            l_skip,
198            height: 1 << log_stacked_height,
199            width: stacked_width,
200            sorted_cols,
201            mat_starts,
202        })
203    }
204
205    /// Raw unsafe constructor
206    pub fn from_raw_parts(
207        l_skip: usize,
208        log_stacked_height: usize,
209        sorted_cols: Vec<(usize, usize, StackedSlice)>,
210    ) -> Result<Self, StackedPcsError> {
211        let height = 1 << log_stacked_height;
212        let width = sorted_cols
213            .iter()
214            .map(|(_, _, slice)| slice.col_idx + 1)
215            .max()
216            .unwrap_or(0);
217        let mut mat_starts = Vec::new();
218        for (idx, (mat_idx, _, _)) in sorted_cols.iter().enumerate() {
219            if idx == 0 || *mat_idx + 1 != mat_starts.len() {
220                if *mat_idx != mat_starts.len() {
221                    return Err(StackedPcsError::LayoutRawPartsMatIdx {
222                        mat_idx: *mat_idx,
223                        mat_starts_len: mat_starts.len(),
224                    });
225                }
226                mat_starts.push(idx);
227            }
228        }
229        Ok(Self {
230            l_skip,
231            height,
232            width,
233            sorted_cols,
234            mat_starts,
235        })
236    }
237
238    pub fn unstacked_slices_iter(&self) -> impl Iterator<Item = &StackedSlice> {
239        self.sorted_cols.iter().map(|(_, _, s)| s)
240    }
241
242    /// `(mat_idx, col_idx)` should be indexing into the unstacked collection of matrices.
243    pub fn get(&self, mat_idx: usize, col_idx: usize) -> Option<&StackedSlice> {
244        let idx = self.mat_starts[mat_idx];
245        if idx + col_idx >= self.sorted_cols.len() {
246            return None;
247        }
248        let (mat_idx1, col_idx1, s) = &self.sorted_cols[idx + col_idx];
249        debug_assert_eq!(*mat_idx1, mat_idx);
250        debug_assert_eq!(*col_idx1, col_idx);
251        Some(s)
252    }
253
254    pub fn width_of(&self, mat_idx: usize) -> usize {
255        let start_idx = self.mat_starts[mat_idx];
256        debug_assert_eq!(self.sorted_cols[start_idx].0, mat_idx);
257        debug_assert_eq!(self.sorted_cols[start_idx].1, 0);
258        let next_idx = *self
259            .mat_starts
260            .get(mat_idx + 1)
261            .unwrap_or(&self.sorted_cols.len());
262        debug_assert_ne!(next_idx, usize::MAX);
263        next_idx - start_idx
264    }
265
266    /// Due to the definition of stacking, in a column major matrix the lifted columns of the
267    /// unstacked matrix will always be contiguous in memory within the stacked matrix, so we
268    /// can return the sub-view.
269    pub fn mat_view<'a, F>(
270        &self,
271        unstacked_mat_idx: usize,
272        stacked_matrix: &'a ColMajorMatrix<F>,
273    ) -> StridedColMajorMatrixView<'a, F> {
274        let col_slices = self
275            .sorted_cols
276            .iter()
277            .filter(|(m, _, _)| *m == unstacked_mat_idx)
278            .collect_vec();
279        let width = col_slices.len();
280        let s = &col_slices[0].2;
281        let lifted_height = s.len(self.l_skip);
282        let stride = s.stride(self.l_skip);
283        let start = col_maj_idx(s.row_idx, s.col_idx, stacked_matrix.height());
284        StridedColMajorMatrixView::new(
285            &stacked_matrix.values[start..start + lifted_height * width],
286            width,
287            stride,
288        )
289    }
290}
291
292/// The `traces` **must** already be in height-sorted order.
293#[instrument(skip_all)]
294pub fn stacked_matrix<F: Field>(
295    l_skip: usize,
296    n_stack: usize,
297    traces: &[&ColMajorMatrix<F>],
298) -> Result<(ColMajorMatrix<F>, StackedLayout), StackedPcsError> {
299    let sorted_meta = traces
300        .iter()
301        .map(|trace| {
302            // height cannot be zero:
303            let log_height = log2_strict_usize(trace.height());
304            (trace.width(), log_height)
305        })
306        .collect_vec();
307    let mut layout = StackedLayout::new(l_skip, l_skip + n_stack, sorted_meta)?;
308    let total_cells: usize = traces
309        .iter()
310        .map(|t| t.height().max(1 << l_skip) * t.width())
311        .sum();
312    let height = 1usize << (l_skip + n_stack);
313    let width = total_cells.div_ceil(height);
314
315    let mut q_mat = F::zero_vec(
316        width
317            .checked_mul(height)
318            .ok_or(StackedPcsError::StackedMatrixOverflow)?,
319    );
320    for (mat_idx, j, s) in &mut layout.sorted_cols {
321        let start = s.col_idx * height + s.row_idx;
322        let t_col = traces[*mat_idx].column(*j);
323        debug_assert_eq!(t_col.len(), 1 << s.log_height);
324        if s.log_height >= l_skip {
325            q_mat[start..start + t_col.len()].copy_from_slice(t_col);
326        } else {
327            // t_col height is smaller than 2^l_skip, so we stride
328            let stride = s.stride(l_skip);
329            for (i, val) in t_col.iter().enumerate() {
330                q_mat[start + i * stride] = *val;
331            }
332        }
333    }
334    Ok((ColMajorMatrix::new(q_mat, width), layout))
335}
336
337/// Computes the Reed-Solomon codeword of each column vector of `eval_matrix` where the rate is
338/// `2^{-log_blowup}`. The column vectors are treated as evaluations of a prismalinear extension on
339/// a hyperprism.
340#[instrument(skip_all)]
341pub fn rs_code_matrix<F: TwoAdicField + Ord>(
342    l_skip: usize,
343    log_blowup: usize,
344    eval_matrix: &ColMajorMatrix<F>,
345) -> Result<ColMajorMatrix<F>, StackedPcsError> {
346    let height = eval_matrix.height();
347    let rs_height = height
348        .checked_shl(log_blowup as u32)
349        .ok_or(StackedPcsError::RsCodeShiftOverflow { height, log_blowup })?;
350    let codewords: Vec<_> = eval_matrix
351        .values
352        .par_chunks_exact(height)
353        .map(|column_evals| {
354            // Convert column evaluations on `D × {0,1}^n` directly into the eval-to-coeff RS
355            // coefficient vector, avoiding redundant interpolation work.
356            let mut coeffs = eval_to_coeff_rs_message(l_skip, column_evals);
357
358            // Compute RS codeword on the resulting univariate polynomial in coefficient form.
359            let dft = Radix2DitParallel::default();
360            coeffs.resize(rs_height, F::ZERO);
361            dft.dft(coeffs)
362        })
363        .collect::<Vec<_>>()
364        .concat();
365
366    Ok(ColMajorMatrix::new(codewords, eval_matrix.width()))
367}
368
369impl<F, Digest> MerkleTree<F, Digest> {
370    pub fn query_stride(&self) -> usize {
371        self.digest_layers[0].len()
372    }
373
374    pub fn proof_depth(&self) -> usize {
375        self.digest_layers.len() - 1
376    }
377}
378
379impl<F, Digest: Clone> MerkleTree<F, Digest> {
380    pub fn root(&self) -> Result<Digest, StackedPcsError> {
381        Ok(self
382            .digest_layers
383            .last()
384            .ok_or(StackedPcsError::MerkleTreeNoRoot)?[0]
385            .clone())
386    }
387
388    pub fn query_merkle_proof(&self, query_idx: usize) -> Result<Vec<Digest>, StackedPcsError> {
389        let stride = self.query_stride();
390        if query_idx >= stride {
391            return Err(StackedPcsError::MerkleTreeQueryOutOfBounds {
392                query_idx,
393                query_stride: stride,
394            });
395        }
396
397        let mut idx = query_idx;
398        let mut proof = Vec::with_capacity(self.proof_depth());
399        for layer in self.digest_layers.iter().take(self.proof_depth()) {
400            let sibling = layer[idx ^ 1].clone();
401            proof.push(sibling);
402            idx >>= 1;
403        }
404        Ok(proof)
405    }
406}
407
408impl<EF: Field, Digest> MerkleTree<EF, Digest>
409where
410    Digest: Copy + Send + Sync,
411{
412    #[instrument(name = "merkle_tree", skip_all)]
413    pub fn new<H: MerkleHasher<Digest = Digest>>(
414        hasher: &H,
415        matrix: ColMajorMatrix<EF>,
416        rows_per_query: usize,
417    ) -> Result<Self, StackedPcsError>
418    where
419        EF: ExtensionField<H::F>,
420    {
421        let height = matrix.height();
422        if height == 0 {
423            return Err(StackedPcsError::MerkleTreeEmptyMatrix);
424        }
425        if !rows_per_query.is_power_of_two() {
426            return Err(StackedPcsError::MerkleTreeRowsPerQueryNotPow2 { rows_per_query });
427        }
428        let num_leaves = height.next_power_of_two();
429        if rows_per_query > num_leaves {
430            return Err(StackedPcsError::MerkleTreeRowsPerQueryExceeded {
431                rows_per_query,
432                num_leaves,
433            });
434        }
435        let row_hashes: Vec<_> = (0..num_leaves)
436            .into_par_iter()
437            .map(|r| {
438                let hash_input: Vec<H::F> = Self::row_iter(&matrix, r)
439                    .flat_map(|ef| ef.as_basis_coefficients_slice().to_vec())
440                    .collect();
441                hasher.hash_slice(&hash_input)
442            })
443            .collect();
444
445        let query_stride = num_leaves / rows_per_query;
446        let mut query_digest_layer = row_hashes;
447        // For the first log2(rows_per_query) layers, we hash in `query_stride` pairs and don't
448        // need to store the digest layers
449        for _ in 0..log2_strict_usize(rows_per_query) {
450            let prev_layer = query_digest_layer;
451            query_digest_layer = (0..prev_layer.len() / 2)
452                .into_par_iter()
453                .map(|i| {
454                    let x = i / query_stride;
455                    let y = i % query_stride;
456                    let left = prev_layer[2 * x * query_stride + y];
457                    let right = prev_layer[(2 * x + 1) * query_stride + y];
458                    hasher.compress(left, right)
459                })
460                .collect();
461        }
462
463        let mut digest_layers = vec![query_digest_layer];
464        while digest_layers
465            .last()
466            .ok_or(StackedPcsError::MerkleTreeNoRoot)?
467            .len()
468            > 1
469        {
470            let prev_layer = digest_layers
471                .last()
472                .ok_or(StackedPcsError::MerkleTreeNoRoot)?;
473            let layer: Vec<_> = prev_layer
474                .par_chunks_exact(2)
475                .map(|pair| hasher.compress(pair[0], pair[1]))
476                .collect();
477            digest_layers.push(layer);
478        }
479
480        Ok(Self {
481            backing_matrix: matrix,
482            digest_layers,
483            rows_per_query,
484        })
485    }
486
487    /// Construct a `MerkleTree` from pre-computed parts without validation.
488    ///
489    /// # Safety
490    ///
491    /// The caller must guarantee:
492    /// - `digest_layers` form a valid Merkle tree over `backing_matrix`: the leaf layer contains
493    ///   correct hashes of the matrix rows and each subsequent layer contains correct compressions
494    ///   of consecutive pairs from the previous layer, terminating in a single root digest.
495    /// - `rows_per_query` is a power of two and does not exceed the number of leaves (i.e.,
496    ///   `backing_matrix.height().next_power_of_two()`).
497    /// - The leaf layer length equals `backing_matrix.height().next_power_of_two() /
498    ///   rows_per_query`.
499    ///
500    /// Violating these invariants will produce incorrect Merkle proofs or panics
501    /// in downstream query/verification code.
502    pub unsafe fn from_raw_parts(
503        backing_matrix: ColMajorMatrix<EF>,
504        digest_layers: Vec<Vec<Digest>>,
505        rows_per_query: usize,
506    ) -> Self {
507        Self {
508            backing_matrix,
509            digest_layers,
510            rows_per_query,
511        }
512    }
513
514    /// Returns the ordered set of opened rows for the given query index.
515    /// The rows are { query_idx + t * query_stride() } for t in 0..rows_per_query.
516    pub fn get_opened_rows(&self, index: usize) -> Result<Vec<Vec<EF>>, StackedPcsError> {
517        let query_stride = self.query_stride();
518        if index >= query_stride {
519            return Err(StackedPcsError::MerkleTreeOpenedRowsOutOfBounds {
520                index,
521                query_stride,
522            });
523        }
524
525        let rows_per_query = self.rows_per_query;
526        let width = self.backing_matrix.width();
527        let mut preimage = Vec::with_capacity(rows_per_query);
528        for row_offset in 0..rows_per_query {
529            let row_idx = row_offset * query_stride + index;
530            let row = Self::row_iter(&self.backing_matrix, row_idx).collect_vec();
531            debug_assert_eq!(
532                row.len(),
533                width,
534                "row width mismatch: expected {width}, got {}",
535                row.len()
536            );
537            preimage.push(row);
538        }
539        Ok(preimage)
540    }
541
542    fn row_iter(matrix: &ColMajorMatrix<EF>, index: usize) -> impl Iterator<Item = EF> + '_ {
543        (0..matrix.width()).map(move |c| matrix.get(index, c).copied().unwrap_or(EF::ZERO))
544    }
545}
546
547#[cfg(test)]
548mod tests {
549    use itertools::Itertools;
550    use openvm_stark_sdk::config::baby_bear_poseidon2::*;
551    use p3_field::PrimeCharacteristicRing;
552
553    use super::*;
554    use crate::prover::ColMajorMatrix;
555
556    #[test]
557    fn test_stacked_matrix_manual_0() {
558        let columns = [vec![1, 2, 3, 4], vec![5, 6], vec![7]]
559            .map(|v| v.into_iter().map(F::from_u32).collect_vec());
560        let mats = columns
561            .into_iter()
562            .map(|c| ColMajorMatrix::new(c, 1))
563            .collect_vec();
564        let mat_refs = mats.iter().collect_vec();
565        let (stacked_mat, layout) = stacked_matrix(0, 2, &mat_refs).unwrap();
566        assert_eq!(stacked_mat.height(), 4);
567        assert_eq!(stacked_mat.width(), 2);
568        assert_eq!(
569            stacked_mat.values,
570            [1, 2, 3, 4, 5, 6, 7, 0].map(F::from_u32).to_vec()
571        );
572        assert_eq!(layout.mat_starts, vec![0, 1, 2]);
573    }
574
575    #[test]
576    fn test_stacked_matrix_manual_strided_0() {
577        let columns = [vec![1, 2, 3, 4], vec![5, 6], vec![7]]
578            .map(|v| v.into_iter().map(F::from_u32).collect_vec());
579        let mats = columns
580            .into_iter()
581            .map(|c| ColMajorMatrix::new(c, 1))
582            .collect_vec();
583        let mat_refs = mats.iter().collect_vec();
584        let (stacked_mat, _layout) = stacked_matrix(2, 0, &mat_refs).unwrap();
585        assert_eq!(stacked_mat.height(), 4);
586        assert_eq!(stacked_mat.width(), 3);
587        assert_eq!(
588            stacked_mat.values,
589            [1, 2, 3, 4, 5, 0, 6, 0, 7, 0, 0, 0]
590                .map(F::from_u32)
591                .to_vec()
592        );
593    }
594
595    #[test]
596    fn test_stacked_matrix_manual_strided_1() {
597        let columns = [vec![1, 2, 3, 4], vec![5, 6], vec![7]]
598            .map(|v| v.into_iter().map(F::from_u32).collect_vec());
599        let mats = columns
600            .into_iter()
601            .map(|c| ColMajorMatrix::new(c, 1))
602            .collect_vec();
603        let mat_refs = mats.iter().collect_vec();
604        let (stacked_mat, _layout) = stacked_matrix(3, 0, &mat_refs).unwrap();
605        assert_eq!(stacked_mat.height(), 8);
606        assert_eq!(stacked_mat.width(), 3);
607        assert_eq!(
608            stacked_mat.values,
609            [
610                [1, 0, 2, 0, 3, 0, 4, 0],
611                [5, 0, 0, 0, 6, 0, 0, 0],
612                [7, 0, 0, 0, 0, 0, 0, 0]
613            ]
614            .into_iter()
615            .flatten()
616            .map(F::from_u32)
617            .collect_vec()
618        );
619    }
620}
openvm_stark_backend/prover/stacked_pcs.rs

openvm_stark_backend/prover/
stacked_pcs.rs
