Skip to main content

sp1_hypercube/air/
sub_builder.rs

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