p3_air/
air.rs

1use core::ops::{Add, Mul, Sub};
2
3use p3_field::{Algebra, ExtensionField, Field, PrimeCharacteristicRing};
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 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: Algebra<Self::F> + Algebra<Self::Var>;
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    /// Assert that the given element is zero.
97    ///
98    /// Where possible, batching multiple assert_zero calls
99    /// into a single assert_zeros call will improve performance.
100    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I);
101
102    /// Assert that every element of a given array is 0.
103    ///
104    /// This should be preferred over calling `assert_zero` multiple times.
105    fn assert_zeros<const N: usize, I: Into<Self::Expr>>(&mut self, array: [I; N]) {
106        for elem in array {
107            self.assert_zero(elem);
108        }
109    }
110
111    /// Assert that a given array consists of only boolean values.
112    fn assert_bools<const N: usize, I: Into<Self::Expr>>(&mut self, array: [I; N]) {
113        let zero_array = array.map(|x| x.into().bool_check());
114        self.assert_zeros(zero_array);
115    }
116
117    /// Assert that `x` element is equal to `1`.
118    fn assert_one<I: Into<Self::Expr>>(&mut self, x: I) {
119        self.assert_zero(x.into() - Self::Expr::ONE);
120    }
121
122    /// Assert that the given elements are equal.
123    fn assert_eq<I1: Into<Self::Expr>, I2: Into<Self::Expr>>(&mut self, x: I1, y: I2) {
124        self.assert_zero(x.into() - y.into());
125    }
126
127    /// Assert that `x` is a boolean, i.e. either `0` or `1`.
128    ///
129    /// Where possible, batching multiple assert_bool calls
130    /// into a single assert_bools call will improve performance.
131    fn assert_bool<I: Into<Self::Expr>>(&mut self, x: I) {
132        self.assert_zero(x.into().bool_check());
133    }
134}
135
136pub trait AirBuilderWithPublicValues: AirBuilder {
137    type PublicVar: Into<Self::Expr> + Copy;
138
139    fn public_values(&self) -> &[Self::PublicVar];
140}
141
142pub trait PairBuilder: AirBuilder {
143    fn preprocessed(&self) -> Self::M;
144}
145
146pub trait ExtensionBuilder: AirBuilder {
147    type EF: ExtensionField<Self::F>;
148
149    type ExprEF: Algebra<Self::Expr> + Algebra<Self::EF>;
150
151    type VarEF: Into<Self::ExprEF> + Copy + Send + Sync;
152
153    fn assert_zero_ext<I>(&mut self, x: I)
154    where
155        I: Into<Self::ExprEF>;
156
157    fn assert_eq_ext<I1, I2>(&mut self, x: I1, y: I2)
158    where
159        I1: Into<Self::ExprEF>,
160        I2: Into<Self::ExprEF>,
161    {
162        self.assert_zero_ext(x.into() - y.into());
163    }
164
165    fn assert_one_ext<I>(&mut self, x: I)
166    where
167        I: Into<Self::ExprEF>,
168    {
169        self.assert_eq_ext(x, Self::ExprEF::ONE)
170    }
171}
172
173pub trait PermutationAirBuilder: ExtensionBuilder {
174    type MP: Matrix<Self::VarEF>;
175
176    type RandomVar: Into<Self::ExprEF> + Copy;
177
178    fn permutation(&self) -> Self::MP;
179
180    fn permutation_randomness(&self) -> &[Self::RandomVar];
181}
182
183#[derive(Debug)]
184pub struct FilteredAirBuilder<'a, AB: AirBuilder> {
185    pub inner: &'a mut AB,
186    condition: AB::Expr,
187}
188
189impl<AB: AirBuilder> FilteredAirBuilder<'_, AB> {
190    pub fn condition(&self) -> AB::Expr {
191        self.condition.clone()
192    }
193}
194
195impl<AB: AirBuilder> AirBuilder for FilteredAirBuilder<'_, AB> {
196    type F = AB::F;
197    type Expr = AB::Expr;
198    type Var = AB::Var;
199    type M = AB::M;
200
201    fn main(&self) -> Self::M {
202        self.inner.main()
203    }
204
205    fn is_first_row(&self) -> Self::Expr {
206        self.inner.is_first_row()
207    }
208
209    fn is_last_row(&self) -> Self::Expr {
210        self.inner.is_last_row()
211    }
212
213    fn is_transition_window(&self, size: usize) -> Self::Expr {
214        self.inner.is_transition_window(size)
215    }
216
217    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
218        self.inner.assert_zero(self.condition() * x.into());
219    }
220}
221
222impl<AB: ExtensionBuilder> ExtensionBuilder for FilteredAirBuilder<'_, AB> {
223    type EF = AB::EF;
224    type ExprEF = AB::ExprEF;
225    type VarEF = AB::VarEF;
226
227    fn assert_zero_ext<I>(&mut self, x: I)
228    where
229        I: Into<Self::ExprEF>,
230    {
231        self.inner.assert_zero_ext(x.into() * self.condition());
232    }
233}
234
235impl<AB: PermutationAirBuilder> PermutationAirBuilder for FilteredAirBuilder<'_, AB> {
236    type MP = AB::MP;
237
238    type RandomVar = AB::RandomVar;
239
240    fn permutation(&self) -> Self::MP {
241        self.inner.permutation()
242    }
243
244    fn permutation_randomness(&self) -> &[Self::RandomVar] {
245        self.inner.permutation_randomness()
246    }
247}
p3_air/air.rs

p3_air/
air.rs
