p3_air/
air.rs

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

p3_air/
air.rs
