openvm_recursion_circuit/
utils.rs

1use std::ops::Index;
2
3use openvm_poseidon2_air::POSEIDON2_WIDTH;
4use openvm_stark_backend::{interaction::Interaction, FiatShamirTranscript};
5use openvm_stark_sdk::config::baby_bear_poseidon2::{
6    poseidon2_perm, BabyBearPoseidon2Config, CHUNK, D_EF, F,
7};
8use p3_air::AirBuilder;
9use p3_field::{extension::BinomiallyExtendable, Field, PrimeCharacteristicRing};
10use p3_symmetric::Permutation;
11
12/// Returns the number of transcript slots consumed by a proof-of-work check.
13///
14/// When `pow_bits > 0`, PoW uses 2 transcript slots (1 observe + 1 sample).
15/// When `pow_bits == 0`, the stark-backend's `check_witness`/`grind` skip
16/// observe/sample entirely, consuming 0 slots.
17#[inline]
18pub const fn pow_tidx_count(pow_bits: usize) -> usize {
19    if pow_bits > 0 {
20        2
21    } else {
22        0
23    }
24}
25
26/// Runs the PoW observe/sample in a preflight transcript, returning the sample
27/// (or `F::ZERO` when `pow_bits == 0`).
28pub fn pow_observe_sample(
29    ts: &mut impl FiatShamirTranscript<BabyBearPoseidon2Config>,
30    pow_bits: usize,
31    witness: F,
32) -> F {
33    if pow_bits > 0 {
34        ts.observe(witness);
35        ts.sample()
36    } else {
37        F::ZERO
38    }
39}
40
41pub fn base_to_ext<FA>(x: impl Into<FA>) -> [FA; D_EF]
42where
43    FA: PrimeCharacteristicRing,
44{
45    [x.into(), FA::ZERO, FA::ZERO, FA::ZERO]
46}
47
48pub fn ext_field_one_minus<FA>(x: [impl Into<FA>; D_EF]) -> [FA; D_EF]
49where
50    FA: PrimeCharacteristicRing,
51{
52    let [x0, x1, x2, x3] = x.map(Into::into);
53    [FA::ONE - x0, -x1, -x2, -x3]
54}
55
56pub fn ext_field_add<FA>(x: [impl Into<FA>; D_EF], y: [impl Into<FA>; D_EF]) -> [FA; D_EF]
57where
58    FA: PrimeCharacteristicRing,
59{
60    let [x0, x1, x2, x3] = x.map(Into::into);
61    let [y0, y1, y2, y3] = y.map(Into::into);
62    [x0 + y0, x1 + y1, x2 + y2, x3 + y3]
63}
64
65pub fn ext_field_subtract<FA>(x: [impl Into<FA>; D_EF], y: [impl Into<FA>; D_EF]) -> [FA; D_EF]
66where
67    FA: PrimeCharacteristicRing,
68{
69    let [x0, x1, x2, x3] = x.map(Into::into);
70    let [y0, y1, y2, y3] = y.map(Into::into);
71    [x0 - y0, x1 - y1, x2 - y2, x3 - y3]
72}
73
74pub fn ext_field_multiply<FA>(x: [impl Into<FA>; D_EF], y: [impl Into<FA>; D_EF]) -> [FA; D_EF]
75where
76    FA: PrimeCharacteristicRing,
77    FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
78{
79    let [x0, x1, x2, x3] = x.map(Into::into);
80    let [y0, y1, y2, y3] = y.map(Into::into);
81
82    let w = FA::from_prime_subfield(FA::PrimeSubfield::W);
83
84    let z0_beta_terms = x1.clone() * y3.clone() + x2.clone() * y2.clone() + x3.clone() * y1.clone();
85    let z1_beta_terms = x2.clone() * y3.clone() + x3.clone() * y2.clone();
86    let z2_beta_terms = x3.clone() * y3.clone();
87
88    [
89        x0.clone() * y0.clone() + z0_beta_terms * w.clone(),
90        x0.clone() * y1.clone() + x1.clone() * y0.clone() + z1_beta_terms * w.clone(),
91        x0.clone() * y2.clone()
92            + x1.clone() * y1.clone()
93            + x2.clone() * y0.clone()
94            + z2_beta_terms * w,
95        x0 * y3 + x1 * y2 + x2 * y1 + x3 * y0,
96    ]
97}
98
99pub fn ext_field_add_scalar<FA>(x: [impl Into<FA>; D_EF], y: impl Into<FA>) -> [FA; D_EF]
100where
101    FA: PrimeCharacteristicRing,
102{
103    let [x0, x1, x2, x3] = x.map(Into::into);
104    [x0 + y.into(), x1, x2, x3]
105}
106
107pub fn ext_field_subtract_scalar<FA>(x: [impl Into<FA>; D_EF], y: impl Into<FA>) -> [FA; D_EF]
108where
109    FA: PrimeCharacteristicRing,
110{
111    let [x0, x1, x2, x3] = x.map(Into::into);
112    [x0 - y.into(), x1, x2, x3]
113}
114
115pub fn scalar_subtract_ext_field<FA>(x: impl Into<FA>, y: [impl Into<FA>; D_EF]) -> [FA; D_EF]
116where
117    FA: PrimeCharacteristicRing,
118{
119    let [y0, y1, y2, y3] = y.map(Into::into);
120    [x.into() - y0, -y1, -y2, -y3]
121}
122
123pub fn ext_field_multiply_scalar<FA>(x: [impl Into<FA>; D_EF], y: impl Into<FA>) -> [FA; D_EF]
124where
125    FA: PrimeCharacteristicRing,
126{
127    let [x0, x1, x2, x3] = x.map(Into::into);
128    let y = y.into();
129    [x0 * y.clone(), x1 * y.clone(), x2 * y.clone(), x3 * y]
130}
131
132pub fn eq_1<FA>(x: [impl Into<FA>; D_EF], y: [impl Into<FA>; D_EF]) -> [FA; D_EF]
133where
134    FA: PrimeCharacteristicRing,
135    FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
136{
137    let x = x.map(Into::into);
138    let y = y.map(Into::into);
139
140    let xy = ext_field_multiply::<FA>(x.clone(), y.clone());
141    let two_xy = ext_field_multiply_scalar::<FA>(xy, FA::TWO);
142    let x_plus_y = ext_field_add::<FA>(x, y);
143    let one_minus_x_plus_y = scalar_subtract_ext_field::<FA>(FA::ONE, x_plus_y);
144    ext_field_add(one_minus_x_plus_y, two_xy)
145}
146
147/// Per-coordinate Möbius-adjusted equality kernel for eval-to-coeff RS encoding:
148/// ```text
149/// mobius_eq_1(u, x) = (1 - 2*u)*(1 - x) + u*x
150///                   = 1 - x - 2*u + 3*u*x
151/// ```
152pub fn mobius_eq_1<FA>(u: [impl Into<FA>; D_EF], x: [impl Into<FA>; D_EF]) -> [FA; D_EF]
153where
154    FA: PrimeCharacteristicRing,
155    FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
156{
157    let omega = u.map(Into::into);
158    let x = x.map(Into::into);
159    // mobius_eq_1(u, x) = 1 - x - 2*u + 3*u*x
160    let omega_x = ext_field_multiply::<FA>(omega.clone(), x.clone());
161    let three_omega_x = ext_field_multiply_scalar::<FA>(omega_x, FA::from_u8(3));
162    let two_omega = ext_field_multiply_scalar::<FA>(omega.clone(), FA::TWO);
163    let x_plus_2omega = ext_field_add::<FA>(x, two_omega);
164    let one_minus_rest = scalar_subtract_ext_field::<FA>(FA::ONE, x_plus_2omega);
165    ext_field_add(one_minus_rest, three_omega_x)
166}
167
168pub fn assert_zeros<AB, const N: usize>(builder: &mut AB, array: [impl Into<AB::Expr>; N])
169where
170    AB: AirBuilder,
171{
172    for elem in array.into_iter() {
173        builder.assert_zero(elem);
174    }
175}
176
177pub fn assert_one_ext<AB>(builder: &mut AB, array: [impl Into<AB::Expr>; D_EF])
178where
179    AB: AirBuilder,
180{
181    for (i, elem) in array.into_iter().enumerate() {
182        if i == 0 {
183            builder.assert_one(elem);
184        } else {
185            builder.assert_zero(elem);
186        }
187    }
188}
189
190pub trait FlattenedLayout {
191    type Index;
192
193    fn len(&self) -> usize;
194    fn is_empty(&self) -> bool {
195        self.len() == 0
196    }
197    fn offset(&self, idx: Self::Index) -> usize;
198}
199
200#[derive(Debug, Clone)]
201pub struct FlattenedVec<T, L> {
202    data: Vec<T>,
203    layout: L,
204}
205
206impl<T, L: FlattenedLayout> FlattenedVec<T, L> {
207    pub fn from_parts(layout: L, data: Vec<T>) -> Self {
208        let layout_len = layout.len();
209        assert_eq!(
210            data.len(),
211            layout_len,
212            "flattened vec layout/data length mismatch: data_len={}, layout_len={layout_len}",
213            data.len(),
214        );
215        Self { data, layout }
216    }
217
218    pub fn len(&self) -> usize {
219        self.data.len()
220    }
221
222    pub fn is_empty(&self) -> bool {
223        self.data.is_empty()
224    }
225
226    pub fn layout(&self) -> &L {
227        &self.layout
228    }
229
230    pub fn as_slice(&self) -> &[T] {
231        &self.data
232    }
233
234    pub fn as_mut_slice(&mut self) -> &mut [T] {
235        &mut self.data
236    }
237}
238
239impl<T, L: FlattenedLayout> Index<L::Index> for FlattenedVec<T, L> {
240    type Output = T;
241
242    fn index(&self, index: L::Index) -> &Self::Output {
243        &self.data[self.layout.offset(index)]
244    }
245}
246
247#[derive(Debug, Clone)]
248pub(crate) struct MultiProofVecVec<T> {
249    data: Vec<T>,
250    bounds: Vec<usize>,
251}
252
253impl<T> MultiProofVecVec<T> {
254    pub(crate) fn new() -> Self {
255        Self {
256            data: Vec::new(),
257            bounds: vec![0],
258        }
259    }
260
261    pub(crate) fn with_capacity(capacity: usize) -> Self {
262        Self {
263            data: Vec::with_capacity(capacity),
264            bounds: vec![0],
265        }
266    }
267
268    pub(crate) fn push(&mut self, x: T) {
269        self.data.push(x);
270    }
271
272    pub(crate) fn extend_from_slice(&mut self, slice: &[T])
273    where
274        T: Clone,
275    {
276        self.data.extend_from_slice(slice);
277    }
278
279    pub(crate) fn extend(&mut self, iter: impl IntoIterator<Item = T>) {
280        self.data.extend(iter);
281    }
282
283    pub(crate) fn end_proof(&mut self) {
284        self.bounds.push(self.data.len());
285    }
286
287    pub(crate) fn len(&self) -> usize {
288        self.data.len()
289    }
290
291    pub(crate) fn num_proofs(&self) -> usize {
292        self.bounds.len() - 1
293    }
294
295    pub(crate) fn data(&self) -> &[T] {
296        &self.data
297    }
298}
299
300impl<T> Index<usize> for MultiProofVecVec<T> {
301    type Output = [T];
302
303    fn index(&self, index: usize) -> &Self::Output {
304        debug_assert!(index < self.num_proofs());
305        &self.data[self.bounds[index]..self.bounds[index + 1]]
306    }
307}
308
309#[derive(Debug, Clone)]
310pub struct MultiVecWithBounds<T, const DIM_MINUS_ONE: usize> {
311    pub data: Vec<T>,
312    pub bounds: [Vec<usize>; DIM_MINUS_ONE],
313}
314
315impl<T, const DIM_MINUS_ONE: usize> MultiVecWithBounds<T, DIM_MINUS_ONE> {
316    #[allow(clippy::new_without_default)]
317    pub fn new() -> Self {
318        Self {
319            data: Vec::new(),
320            bounds: core::array::from_fn(|_| vec![0]),
321        }
322    }
323
324    pub fn with_capacity(capacity: usize) -> Self {
325        Self {
326            data: Vec::with_capacity(capacity),
327            bounds: core::array::from_fn(|_| vec![0]),
328        }
329    }
330
331    pub fn push(&mut self, x: T) {
332        self.data.push(x);
333    }
334
335    pub fn close_level(&mut self, level: usize) {
336        debug_assert!(level < DIM_MINUS_ONE);
337        for i in level..DIM_MINUS_ONE - 1 {
338            self.bounds[i].push(self.bounds[i + 1].len());
339        }
340        self.bounds[DIM_MINUS_ONE - 1].push(self.data.len());
341    }
342
343    pub fn extend(&mut self, iter: impl IntoIterator<Item = T>) {
344        self.data.extend(iter);
345    }
346}
347
348impl<T, const DIM_MINUS_ONE: usize> Index<[usize; DIM_MINUS_ONE]>
349    for MultiVecWithBounds<T, DIM_MINUS_ONE>
350{
351    type Output = [T];
352
353    fn index(&self, index: [usize; DIM_MINUS_ONE]) -> &Self::Output {
354        let mut idx = 0;
355        #[allow(clippy::needless_range_loop)]
356        for i in 0..DIM_MINUS_ONE {
357            idx += index[i];
358            if i < DIM_MINUS_ONE - 1 {
359                idx = self.bounds[i][idx];
360            }
361        }
362        &self.data[self.bounds[DIM_MINUS_ONE - 1][idx]..self.bounds[DIM_MINUS_ONE - 1][idx + 1]]
363    }
364}
365
366pub fn interpolate_quadratic<FA>(
367    pre_claim: [impl Into<FA>; D_EF],
368    ev1: [impl Into<FA>; D_EF],
369    ev2: [impl Into<FA>; D_EF],
370    alpha: [impl Into<FA>; D_EF],
371) -> [FA; D_EF]
372where
373    FA: PrimeCharacteristicRing,
374    FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
375{
376    let pre_claim = pre_claim.map(Into::into);
377    let ev1 = ev1.map(Into::into);
378    let ev2 = ev2.map(Into::into);
379    let alpha = alpha.map(Into::into);
380
381    let ev0 = ext_field_subtract::<FA>(pre_claim.clone(), ev1.clone());
382    let s1 = ext_field_subtract::<FA>(ev1.clone(), ev0.clone());
383    let s2 = ext_field_subtract::<FA>(ev2.clone(), ev1.clone());
384    let p = ext_field_multiply::<FA>(
385        ext_field_subtract::<FA>(s2, s1.clone()),
386        base_to_ext::<FA>(FA::from_prime_subfield(FA::PrimeSubfield::TWO.inverse())),
387    );
388    let q = ext_field_subtract::<FA>(s1, p.clone());
389    ext_field_add::<FA>(
390        ev0,
391        ext_field_multiply::<FA>(
392            alpha.clone(),
393            ext_field_add::<FA>(q, ext_field_multiply::<FA>(p, alpha)),
394        ),
395    )
396}
397
398pub fn poseidon2_hash_slice(vals: &[F]) -> ([F; CHUNK], Vec<[F; POSEIDON2_WIDTH]>) {
399    let num_chunks = vals.len().div_ceil(CHUNK);
400    let mut pre_states = Vec::with_capacity(num_chunks);
401    let perm = poseidon2_perm();
402    let mut state = [F::ZERO; POSEIDON2_WIDTH];
403    let mut i = 0;
404    for &val in vals {
405        state[i] = val;
406        i += 1;
407        if i == CHUNK {
408            pre_states.push(state);
409            perm.permute_mut(&mut state);
410            i = 0;
411        }
412    }
413    if i != 0 {
414        pre_states.push(state);
415        perm.permute_mut(&mut state);
416    }
417    (state[..CHUNK].try_into().unwrap(), pre_states)
418}
419
420#[inline]
421pub fn poseidon2_hash_slice_with_states(
422    vals: &[F],
423) -> (
424    [F; CHUNK],
425    Vec<[F; POSEIDON2_WIDTH]>,
426    Vec<[F; POSEIDON2_WIDTH]>,
427) {
428    let num_chunks = vals.len().div_ceil(CHUNK);
429    let mut pre_states = Vec::with_capacity(num_chunks);
430    let mut post_states = Vec::with_capacity(num_chunks);
431    let perm = poseidon2_perm();
432    let mut state = [F::ZERO; POSEIDON2_WIDTH];
433    let mut i = 0;
434    for &val in vals {
435        state[i] = val;
436        i += 1;
437        if i == CHUNK {
438            pre_states.push(state);
439            perm.permute_mut(&mut state);
440            post_states.push(state);
441            i = 0;
442        }
443    }
444    if i != 0 {
445        pre_states.push(state);
446        perm.permute_mut(&mut state);
447        post_states.push(state);
448    }
449    (state[..CHUNK].try_into().unwrap(), pre_states, post_states)
450}
451
452/// The number of fields corresponding to an interaction.
453/// Is defined here because only makes sense in terms of recursion
454/// (e.g. indicates the number of rows somewhere).
455pub(crate) fn interaction_length<T>(interaction: &Interaction<T>) -> usize {
456    interaction.message.len() + 2
457}
openvm_recursion_circuit/utils.rs

openvm_recursion_circuit/
utils.rs
