p3_air/
air.rs

1use core::ops::{Add, Mul, Sub};
2
3use p3_field::{AbstractExtensionField, AbstractField, ExtensionField, Field};
4use p3_matrix::Matrix;
5use p3_matrix::dense::RowMajorMatrix;
6
7/// An AIR (algebraic intermediate representation).
8pub trait BaseAir<F>: Sync {
9    /// The number of columns (a.k.a. registers) in this AIR.
10    fn width(&self) -> usize;
11
12    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<F>> {
13        None
14    }
15}
16
17/// An AIR that works with a particular `AirBuilder`.
18pub trait Air<AB: AirBuilder>: BaseAir<AB::F> {
19    fn eval(&self, builder: &mut AB);
20}
21
22pub trait AirBuilder: Sized {
23    type F: Field;
24
25    type Expr: AbstractField
26        + From<Self::F>
27        + Add<Self::Var, Output = Self::Expr>
28        + Add<Self::F, Output = Self::Expr>
29        + Sub<Self::Var, Output = Self::Expr>
30        + Sub<Self::F, Output = Self::Expr>
31        + Mul<Self::Var, Output = Self::Expr>
32        + Mul<Self::F, Output = Self::Expr>;
33
34    type Var: Into<Self::Expr>
35        + Copy
36        + Send
37        + Sync
38        + Add<Self::F, Output = Self::Expr>
39        + Add<Self::Var, Output = Self::Expr>
40        + Add<Self::Expr, Output = Self::Expr>
41        + Sub<Self::F, Output = Self::Expr>
42        + Sub<Self::Var, Output = Self::Expr>
43        + Sub<Self::Expr, Output = Self::Expr>
44        + Mul<Self::F, Output = Self::Expr>
45        + Mul<Self::Var, Output = Self::Expr>
46        + Mul<Self::Expr, Output = Self::Expr>;
47
48    type M: Matrix<Self::Var>;
49
50    fn main(&self) -> Self::M;
51
52    fn is_first_row(&self) -> Self::Expr;
53    fn is_last_row(&self) -> Self::Expr;
54    fn is_transition(&self) -> Self::Expr {
55        self.is_transition_window(2)
56    }
57    fn is_transition_window(&self, size: usize) -> Self::Expr;
58
59    /// Returns a sub-builder whose constraints are enforced only when `condition` is nonzero.
60    fn when<I: Into<Self::Expr>>(&mut self, condition: I) -> FilteredAirBuilder<'_, Self> {
61        FilteredAirBuilder {
62            inner: self,
63            condition: condition.into(),
64        }
65    }
66
67    /// Returns a sub-builder whose constraints are enforced only when `x != y`.
68    fn when_ne<I1: Into<Self::Expr>, I2: Into<Self::Expr>>(
69        &mut self,
70        x: I1,
71        y: I2,
72    ) -> FilteredAirBuilder<'_, Self> {
73        self.when(x.into() - y.into())
74    }
75
76    /// Returns a sub-builder whose constraints are enforced only on the first row.
77    fn when_first_row(&mut self) -> FilteredAirBuilder<'_, Self> {
78        self.when(self.is_first_row())
79    }
80
81    /// Returns a sub-builder whose constraints are enforced only on the last row.
82    fn when_last_row(&mut self) -> FilteredAirBuilder<'_, Self> {
83        self.when(self.is_last_row())
84    }
85
86    /// Returns a sub-builder whose constraints are enforced on all rows except the last.
87    fn when_transition(&mut self) -> FilteredAirBuilder<'_, Self> {
88        self.when(self.is_transition())
89    }
90
91    /// Returns a sub-builder whose constraints are enforced on all rows except the last `size - 1`.
92    fn when_transition_window(&mut self, size: usize) -> FilteredAirBuilder<'_, Self> {
93        self.when(self.is_transition_window(size))
94    }
95
96    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I);
97
98    fn assert_one<I: Into<Self::Expr>>(&mut self, x: I) {
99        self.assert_zero(x.into() - Self::Expr::one());
100    }
101
102    fn assert_eq<I1: Into<Self::Expr>, I2: Into<Self::Expr>>(&mut self, x: I1, y: I2) {
103        self.assert_zero(x.into() - y.into());
104    }
105
106    /// Assert that `x` is a boolean, i.e. either 0 or 1.
107    fn assert_bool<I: Into<Self::Expr>>(&mut self, x: I) {
108        let x = x.into();
109        self.assert_zero(x.clone() * (x - Self::Expr::one()));
110    }
111}
112
113pub trait AirBuilderWithPublicValues: AirBuilder {
114    type PublicVar: Into<Self::Expr> + Copy;
115
116    fn public_values(&self) -> &[Self::PublicVar];
117}
118
119pub trait PairBuilder: AirBuilder {
120    fn preprocessed(&self) -> Self::M;
121}
122
123pub trait ExtensionBuilder: AirBuilder {
124    type EF: ExtensionField<Self::F>;
125
126    type ExprEF: AbstractExtensionField<Self::Expr, F = Self::EF>
127        + Mul<Self::VarEF, Output = Self::ExprEF>;
128
129    type VarEF: Into<Self::ExprEF> + Copy + Send + Sync;
130
131    fn assert_zero_ext<I>(&mut self, x: I)
132    where
133        I: Into<Self::ExprEF>;
134
135    fn assert_eq_ext<I1, I2>(&mut self, x: I1, y: I2)
136    where
137        I1: Into<Self::ExprEF>,
138        I2: Into<Self::ExprEF>,
139    {
140        self.assert_zero_ext(x.into() - y.into());
141    }
142
143    fn assert_one_ext<I>(&mut self, x: I)
144    where
145        I: Into<Self::ExprEF>,
146    {
147        self.assert_eq_ext(x, Self::ExprEF::one())
148    }
149}
150
151pub trait PermutationAirBuilder: ExtensionBuilder {
152    type MP: Matrix<Self::VarEF>;
153
154    type RandomVar: Into<Self::ExprEF> + Copy;
155
156    fn permutation(&self) -> Self::MP;
157
158    fn permutation_randomness(&self) -> &[Self::RandomVar];
159}
160
161#[derive(Debug)]
162pub struct FilteredAirBuilder<'a, AB: AirBuilder> {
163    pub inner: &'a mut AB,
164    condition: AB::Expr,
165}
166
167impl<'a, AB: AirBuilder> AirBuilder for FilteredAirBuilder<'a, AB> {
168    type F = AB::F;
169    type Expr = AB::Expr;
170    type Var = AB::Var;
171    type M = AB::M;
172
173    fn main(&self) -> Self::M {
174        self.inner.main()
175    }
176
177    fn is_first_row(&self) -> Self::Expr {
178        self.inner.is_first_row()
179    }
180
181    fn is_last_row(&self) -> Self::Expr {
182        self.inner.is_last_row()
183    }
184
185    fn is_transition_window(&self, size: usize) -> Self::Expr {
186        self.inner.is_transition_window(size)
187    }
188
189    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
190        self.inner.assert_zero(self.condition.clone() * x.into());
191    }
192}
193
194impl<'a, AB: ExtensionBuilder> ExtensionBuilder for FilteredAirBuilder<'a, AB> {
195    type EF = AB::EF;
196    type ExprEF = AB::ExprEF;
197    type VarEF = AB::VarEF;
198
199    fn assert_zero_ext<I>(&mut self, x: I)
200    where
201        I: Into<Self::ExprEF>,
202    {
203        self.inner
204            .assert_zero_ext(x.into() * self.condition.clone());
205    }
206}
207
208impl<'a, AB: PermutationAirBuilder> PermutationAirBuilder for FilteredAirBuilder<'a, AB> {
209    type MP = AB::MP;
210
211    type RandomVar = AB::RandomVar;
212
213    fn permutation(&self) -> Self::MP {
214        self.inner.permutation()
215    }
216
217    fn permutation_randomness(&self) -> &[Self::RandomVar] {
218        self.inner.permutation_randomness()
219    }
220}
221
222impl<'a, AB: AirBuilderWithPublicValues> AirBuilderWithPublicValues for FilteredAirBuilder<'a, AB> {
223    type PublicVar = AB::PublicVar;
224
225    fn public_values(&self) -> &[Self::PublicVar] {
226        self.inner.public_values()
227    }
228}
p3_air/air.rs

p3_air/
air.rs
