openvm_stark_backend/air_builders/
sub.rs

Help
1// Copied from sp1/core under MIT license
2
3use std::{
4    iter::{Skip, Take},
5    ops::{Deref, Range},
6};
7
8use p3_air::{AirBuilder, BaseAir};
9use p3_matrix::Matrix;
10
11/// A submatrix of a matrix.  The matrix will contain a subset of the columns of `self.inner`.
12pub struct SubMatrixRowSlices<M: Matrix<T>, T: Send + Sync> {
13    inner: M,
14    column_range: Range<usize>,
15    _phantom: std::marker::PhantomData<T>,
16}
17
18impl<M: Matrix<T>, T: Send + Sync> SubMatrixRowSlices<M, T> {
19    pub const fn new(inner: M, column_range: Range<usize>) -> Self {
20        Self {
21            inner,
22            column_range,
23            _phantom: std::marker::PhantomData,
24        }
25    }
26}
27
28/// Implement `Matrix` for `SubMatrixRowSlices`.
29impl<M: Matrix<T>, T: Send + Sync> Matrix<T> for SubMatrixRowSlices<M, T> {
30    type Row<'a>
31        = Skip<Take<M::Row<'a>>>
32    where
33        Self: 'a;
34
35    #[inline]
36    fn row(&self, r: usize) -> Self::Row<'_> {
37        self.inner
38            .row(r)
39            .take(self.column_range.end)
40            .skip(self.column_range.start)
41    }
42
43    #[inline]
44    fn row_slice(&self, r: usize) -> impl Deref<Target = [T]> {
45        self.row(r).collect::<Vec<_>>()
46    }
47
48    #[inline]
49    fn width(&self) -> usize {
50        self.column_range.len()
51    }
52
53    #[inline]
54    fn height(&self) -> usize {
55        self.inner.height()
56    }
57}
58
59/// A builder used to eval a sub-air.  This will handle enforcing constraints for a subset of a
60/// trace matrix.  E.g. if a particular air needs to be enforced for a subset of the columns of
61/// the trace, then the SubAirBuilder can be used.
62pub struct SubAirBuilder<'a, AB: AirBuilder, SubAir: BaseAir<T>, T> {
63    inner: &'a mut AB,
64    column_range: Range<usize>,
65    _phantom: std::marker::PhantomData<(SubAir, T)>,
66}
67
68impl<'a, AB: AirBuilder, SubAir: BaseAir<T>, T> SubAirBuilder<'a, AB, SubAir, T> {
69    pub fn new(inner: &'a mut AB, column_range: Range<usize>) -> Self {
70        Self {
71            inner,
72            column_range,
73            _phantom: std::marker::PhantomData,
74        }
75    }
76}
77
78/// Implement `AirBuilder` for `SubAirBuilder`.
79impl<AB: AirBuilder, SubAir: BaseAir<F>, F> AirBuilder for SubAirBuilder<'_, AB, SubAir, F> {
80    type F = AB::F;
81    type Expr = AB::Expr;
82    type Var = AB::Var;
83    type M = SubMatrixRowSlices<AB::M, Self::Var>;
84
85    fn main(&self) -> Self::M {
86        let matrix = self.inner.main();
87
88        SubMatrixRowSlices::new(matrix, self.column_range.clone())
89    }
90
91    fn is_first_row(&self) -> Self::Expr {
92        self.inner.is_first_row()
93    }
94
95    fn is_last_row(&self) -> Self::Expr {
96        self.inner.is_last_row()
97    }
98
99    fn is_transition_window(&self, size: usize) -> Self::Expr {
100        self.inner.is_transition_window(size)
101    }
102
103    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
104        self.inner.assert_zero(x.into());
105    }
106}
openvm_stark_backend/air_builders/sub.rs

openvm_stark_backend/air_builders/
sub.rs
