openvm_recursion_circuit/batch_constraint/eq_airs/eq_neg/
trace.rs

1use std::borrow::BorrowMut;
2
3use openvm_stark_backend::{
4    keygen::types::MultiStarkVerifyingKey, poly_common::eval_eq_uni_at_one,
5};
6use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, F};
7use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, TwoAdicField};
8use p3_matrix::dense::RowMajorMatrix;
9
10use crate::{
11    batch_constraint::{eq_airs::eq_neg::air::EqNegCols, SelectorCount},
12    system::Preflight,
13    tracegen::RowMajorChip,
14    utils::MultiVecWithBounds,
15};
16
17pub struct EqNegTraceGenerator;
18
19impl RowMajorChip<F> for EqNegTraceGenerator {
20    type Ctx<'a> = (
21        &'a MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
22        &'a [&'a Preflight],
23        &'a MultiVecWithBounds<SelectorCount, 1>,
24    );
25
26    #[tracing::instrument(level = "trace", skip_all)]
27    fn generate_trace(
28        &self,
29        ctx: &Self::Ctx<'_>,
30        required_height: Option<usize>,
31    ) -> Option<RowMajorMatrix<F>> {
32        let (vk, preflights, selector_counts) = ctx;
33        let l_skip = vk.inner.params.l_skip;
34        let width = EqNegCols::<usize>::width();
35
36        if l_skip == 0 {
37            let ret = if let Some(height) = required_height {
38                // We essentially fill the trace with dummy rows as we do below, but
39                // instead of proof_idx being preflights.len() it is 0
40                Some(RowMajorMatrix::new(vec![F::ZERO; height * width], width))
41            } else {
42                Some(RowMajorMatrix::new(vec![], width))
43            };
44            return ret;
45        }
46
47        let total_valid = preflights.len() * (l_skip * (l_skip + 3)) / 2;
48
49        let padded_height = if let Some(height) = required_height {
50            if height < total_valid {
51                return None;
52            }
53            height
54        } else {
55            total_valid.next_power_of_two()
56        };
57        let mut trace = vec![F::ZERO; padded_height * width];
58        let mut chunks = trace.chunks_exact_mut(width);
59
60        for (proof_idx, preflight) in preflights.iter().enumerate() {
61            let initial_omega = F::two_adic_generator(vk.inner.params.l_skip);
62            let initial_u = preflight.stacking.sumcheck_rnd[0];
63            let mut initial_r = preflight.batch_constraint.sumcheck_rnd[0];
64            let mut initial_r_omega = initial_r * initial_omega;
65
66            for neg_hypercube in 0..l_skip {
67                let mut u = initial_u;
68                let mut r = initial_r;
69                let mut r_omega = initial_r_omega;
70
71                let mut prod_u_r = u * (u + r);
72                let mut prod_u_r_omega = u * (u + r_omega);
73                let mut prod_1_r = r + F::ONE;
74                let mut prod_1_r_omega = r_omega + F::ONE;
75
76                let one_half = F::ONE.halve();
77                let mut one_half_pow = one_half;
78
79                for row_idx in 0..=l_skip - neg_hypercube {
80                    let chunk = chunks.next().unwrap();
81                    let cols: &mut EqNegCols<F> = chunk.borrow_mut();
82                    let is_first_hypercube = row_idx == 0;
83                    let is_last_hypercube = row_idx == l_skip - neg_hypercube;
84
85                    cols.proof_idx = F::from_usize(proof_idx);
86                    cols.is_valid = F::ONE;
87                    cols.is_first = F::from_bool(neg_hypercube == 0 && is_first_hypercube);
88                    cols.is_last = F::from_bool(neg_hypercube + 1 == l_skip && is_last_hypercube);
89
90                    cols.neg_hypercube = F::from_usize(neg_hypercube);
91                    cols.neg_hypercube_nz_inv =
92                        cols.neg_hypercube.try_inverse().unwrap_or_default();
93                    cols.row_index = F::from_usize(row_idx);
94                    cols.is_first_hypercube = F::from_bool(is_first_hypercube);
95                    cols.is_last_hypercube = F::from_bool(is_last_hypercube);
96
97                    cols.u_pow.copy_from_slice(u.as_basis_coefficients_slice());
98                    cols.r_pow.copy_from_slice(r.as_basis_coefficients_slice());
99                    cols.r_omega_pow
100                        .copy_from_slice(r_omega.as_basis_coefficients_slice());
101                    u *= u;
102                    r *= r;
103                    r_omega *= r_omega;
104
105                    cols.prod_u_r
106                        .copy_from_slice(prod_u_r.as_basis_coefficients_slice());
107                    cols.prod_u_r_omega
108                        .copy_from_slice(prod_u_r_omega.as_basis_coefficients_slice());
109                    prod_u_r *= u + r;
110                    prod_u_r_omega *= u + r_omega;
111
112                    cols.prod_1_r
113                        .copy_from_slice(prod_1_r.as_basis_coefficients_slice());
114                    cols.prod_1_r_omega
115                        .copy_from_slice(prod_1_r_omega.as_basis_coefficients_slice());
116                    cols.one_half_pow = one_half_pow;
117                    debug_assert_eq!(
118                        prod_1_r * one_half_pow,
119                        eval_eq_uni_at_one(row_idx + 1, initial_r)
120                    );
121                    debug_assert_eq!(
122                        prod_1_r_omega * one_half_pow,
123                        eval_eq_uni_at_one(row_idx + 1, initial_r_omega)
124                    );
125                    prod_1_r *= r + F::ONE;
126                    prod_1_r_omega *= r_omega + F::ONE;
127                    one_half_pow *= one_half;
128
129                    // We only use the count for the last row of `hypercube` or the very first row.
130                    // We use the very first row to hard-code the lookup for log_height = 0 since
131                    // it's disjoint from the other condition.
132                    if neg_hypercube == 0 && row_idx == 0 {
133                        let counts = selector_counts[[proof_idx]][0];
134                        cols.sel_first_count = F::from_usize(counts.first);
135                        cols.sel_last_trans_count = F::from_usize(counts.last + counts.transition);
136                    } else if row_idx == l_skip - neg_hypercube {
137                        let mut counts = selector_counts[[proof_idx]][row_idx];
138                        // On `neg_hypercube`, we collect counts for all heights >= l_skip.
139                        if neg_hypercube == 0 {
140                            for count in &selector_counts[[proof_idx]][l_skip - neg_hypercube + 1..]
141                            {
142                                counts.first += count.first;
143                                counts.last += count.last;
144                                counts.transition += count.transition;
145                            }
146                        }
147                        cols.sel_first_count = F::from_usize(counts.first);
148                        cols.sel_last_trans_count = F::from_usize(counts.last + counts.transition);
149                    }
150                }
151
152                initial_r *= initial_r;
153                initial_r_omega *= initial_r_omega;
154            }
155        }
156
157        for chunk in chunks {
158            let cols: &mut EqNegCols<F> = chunk.borrow_mut();
159            cols.proof_idx = F::from_usize(preflights.len());
160        }
161
162        Some(RowMajorMatrix::new(trace, width))
163    }
164}
openvm_recursion_circuit/batch_constraint/eq_airs/eq_neg/trace.rs

openvm_recursion_circuit/batch_constraint/eq_airs/eq_neg/
trace.rs
