openvm_recursion_circuit/gkr/xi_sampler/
trace.rs

1use core::borrow::BorrowMut;
2
3use openvm_stark_backend::p3_maybe_rayon::prelude::*;
4use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, EF, F};
5use p3_field::{BasedVectorSpace, PrimeCharacteristicRing};
6use p3_matrix::dense::RowMajorMatrix;
7
8use super::GkrXiSamplerCols;
9use crate::tracegen::RowMajorChip;
10
11#[derive(Debug, Clone, Default)]
12pub struct GkrXiSamplerRecord {
13    pub tidx: usize,
14    pub idx: usize,
15    pub xis: Vec<EF>,
16}
17
18pub struct GkrXiSamplerTraceGenerator;
19
20impl RowMajorChip<F> for GkrXiSamplerTraceGenerator {
21    // xi_sampler_records
22    type Ctx<'a> = &'a [GkrXiSamplerRecord];
23
24    #[tracing::instrument(level = "trace", skip_all)]
25    fn generate_trace(
26        &self,
27        ctx: &Self::Ctx<'_>,
28        required_height: Option<usize>,
29    ) -> Option<RowMajorMatrix<F>> {
30        let xi_sampler_records = ctx;
31        let width = GkrXiSamplerCols::<F>::width();
32
33        // Calculate rows per proof (minimum 1 row per proof)
34        let rows_per_proof: Vec<usize> = xi_sampler_records
35            .iter()
36            .map(|record| record.xis.len().max(1))
37            .collect();
38
39        // Calculate total rows
40        let num_valid_rows: usize = rows_per_proof.iter().sum();
41        let height = if let Some(height) = required_height {
42            if height < num_valid_rows {
43                return None;
44            }
45            height
46        } else {
47            num_valid_rows.next_power_of_two()
48        };
49
50        let mut trace = vec![F::ZERO; height * width];
51
52        // Split trace into chunks for each proof
53        let (data_slice, _) = trace.split_at_mut(num_valid_rows * width);
54        let mut trace_slices: Vec<&mut [F]> = Vec::with_capacity(rows_per_proof.len());
55        let mut remaining = data_slice;
56
57        for &num_rows in &rows_per_proof {
58            let chunk_size = num_rows * width;
59            let (chunk, rest) = remaining.split_at_mut(chunk_size);
60            trace_slices.push(chunk);
61            remaining = rest;
62        }
63
64        // Process each proof
65        trace_slices
66            .par_iter_mut()
67            .zip(xi_sampler_records.par_iter())
68            .enumerate()
69            .for_each(|(proof_idx, (proof_trace, xi_sampler_record))| {
70                if xi_sampler_record.xis.is_empty() {
71                    debug_assert_eq!(proof_trace.len(), width);
72                    let row_data = &mut proof_trace[..width];
73                    let cols: &mut GkrXiSamplerCols<F> = row_data.borrow_mut();
74                    cols.is_enabled = F::ONE;
75                    cols.proof_idx = F::from_usize(proof_idx);
76                    cols.is_first_challenge = F::ONE;
77                    cols.is_dummy = F::ONE;
78                    return;
79                }
80
81                let challenge_indices: Vec<usize> = (0..xi_sampler_record.xis.len())
82                    .map(|i| xi_sampler_record.idx + i)
83                    .collect();
84                let tidxs: Vec<usize> = (0..xi_sampler_record.xis.len())
85                    .map(|i| xi_sampler_record.tidx + i * D_EF)
86                    .collect();
87
88                proof_trace
89                    .par_chunks_mut(width)
90                    .zip(
91                        xi_sampler_record
92                            .xis
93                            .par_iter()
94                            .zip(challenge_indices.par_iter())
95                            .zip(tidxs.par_iter()),
96                    )
97                    .enumerate()
98                    .for_each(|(row_idx, (row_data, ((xi, idx), tidx)))| {
99                        let cols: &mut GkrXiSamplerCols<F> = row_data.borrow_mut();
100                        cols.proof_idx = F::from_usize(proof_idx);
101
102                        cols.is_enabled = F::ONE;
103                        cols.is_first_challenge = F::from_bool(row_idx == 0);
104                        cols.tidx = F::from_usize(*tidx);
105                        cols.idx = F::from_usize(*idx);
106                        cols.xi = xi.as_basis_coefficients_slice().try_into().unwrap();
107                    });
108            });
109
110        Some(RowMajorMatrix::new(trace, width))
111    }
112}
openvm_recursion_circuit/gkr/xi_sampler/trace.rs

openvm_recursion_circuit/gkr/xi_sampler/
trace.rs
