p3_batch_stark/
check_constraints.rs

1use p3_air::{
2    AirBuilder, AirBuilderWithPublicValues, ExtensionBuilder, PairBuilder, PermutationAirBuilder,
3};
4use p3_field::{ExtensionField, Field};
5use p3_lookup::lookup_traits::{AirLookupHandler, Lookup, LookupData, LookupGadget};
6use p3_matrix::Matrix;
7use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixView};
8use p3_matrix::stack::{VerticalPair, ViewPair};
9use tracing::instrument;
10
11/// Type alias for the inputs to lookup constraint checking.
12/// - The first element is a slice of [`Lookup`] values (generic over a field `F`) representing the symbolic lookups to be performed.
13/// - The second element is a slice of [`LookupData`] values (generic over an extension field `EF`) representing the lookup data for global lookups.
14/// - The third element is a reference to the [`LookupGadget`] implementation.
15#[allow(unused)]
16type LookupConstraintsInputs<'a, F, EF, LG> = (&'a [Lookup<F>], &'a [LookupData<EF>], &'a LG);
17
18/// Runs constraint checks using a given [AIR](`p3_air::Air`) implementation and trace matrix.
19///
20/// Iterates over every row in `main`, providing both the current and next row
21/// (with wraparound) to the [AIR](`p3_air::Air`) logic. Also injects public values into the builder
22/// for first/last row assertions.
23///
24/// # Arguments
25/// - `air`: The [AIR](`p3_air::Air`) logic to run.
26/// - `main`: The [`RowMajorMatrix`] containing rows of witness values.
27/// - `permutation`: The permutation [`RowMajorMatrix`] (rows of permutation values).
28/// - `permutation_challenges`: The challenges used for the permutation argument.
29/// - `public_values`: Public values provided to the builder.
30/// - `lookup_constraints_inputs`: Inputs necessary to check lookup constraints:
31///     - the symbolic representation of the [`Lookup`] values,
32///     - the [`LookupData`] for global lookups,
33///     - the [`LookupGadget`] implementation.
34#[instrument(name = "check constraints", skip_all)]
35#[allow(unused)]
36pub(crate) fn check_constraints<'b, F, EF, A, LG>(
37    air: &A,
38    main: &RowMajorMatrix<F>,
39    preprocessed: &Option<RowMajorMatrix<F>>,
40    permutation: &RowMajorMatrix<EF>,
41    permutation_challenges: &[EF],
42    public_values: &[F],
43    lookup_constraints_inputs: LookupConstraintsInputs<'b, F, EF, LG>,
44) where
45    F: Field,
46    EF: ExtensionField<F>,
47    A: for<'a> AirLookupHandler<DebugConstraintBuilderWithLookups<'a, F, EF>>,
48    LG: LookupGadget,
49{
50    let height = main.height();
51
52    let (lookups, lookup_data, lookup_gadget) = lookup_constraints_inputs;
53
54    (0..height).for_each(|row_index| {
55        let row_index_next = (row_index + 1) % height;
56
57        // Safety:
58        // - row_index < height so we can use unchecked indexing.
59        // - row_index_next < height so we can use unchecked indexing.
60        let (local, next, prep_local, prep_next, perm_local, perm_next) = unsafe {
61            (
62                main.row_slice_unchecked(row_index),
63                main.row_slice_unchecked(row_index_next),
64                preprocessed
65                    .as_ref()
66                    .map(|p| p.row_slice_unchecked(row_index)),
67                preprocessed
68                    .as_ref()
69                    .map(|p| p.row_slice_unchecked(row_index_next)),
70                permutation.row_slice_unchecked(row_index),
71                permutation.row_slice_unchecked(row_index_next),
72            )
73        };
74        let main = VerticalPair::new(
75            RowMajorMatrixView::new_row(&*local),
76            RowMajorMatrixView::new_row(&*next),
77        );
78
79        let preprocessed_rows_data = prep_local.as_ref().zip(prep_next.as_ref());
80        let preprocessed = preprocessed_rows_data.map(|(prep_local, prep_next)| {
81            VerticalPair::new(
82                RowMajorMatrixView::new_row(&**prep_local),
83                RowMajorMatrixView::new_row(&**prep_next),
84            )
85        });
86
87        let permutation = VerticalPair::new(
88            RowMajorMatrixView::new_row(&*perm_local),
89            RowMajorMatrixView::new_row(&*perm_next),
90        );
91
92        let mut builder = DebugConstraintBuilderWithLookups {
93            row_index,
94            main,
95            preprocessed,
96            permutation,
97            permutation_challenges,
98            public_values,
99            is_first_row: F::from_bool(row_index == 0),
100            is_last_row: F::from_bool(row_index == height - 1),
101            is_transition: F::from_bool(row_index != height - 1),
102        };
103
104        <A as AirLookupHandler<DebugConstraintBuilderWithLookups<'_, F, EF>>>::eval(
105            air,
106            &mut builder,
107            lookups,
108            lookup_data,
109            lookup_gadget,
110        );
111    });
112}
113
114/// A builder that runs constraint assertions during testing.
115///
116/// Used in conjunction with [`check_constraints`] to simulate
117/// an execution trace and verify that the [AIR](`p3_air::Air`) logic enforces all constraints.
118#[derive(Debug)]
119#[allow(unused)]
120pub struct DebugConstraintBuilderWithLookups<'a, F: Field, EF: ExtensionField<F>> {
121    /// The index of the row currently being evaluated.
122    row_index: usize,
123    /// A view of the current and next row as a vertical pair.
124    main: ViewPair<'a, F>,
125    /// A view of the current and next preprocessed row as a vertical pair.
126    preprocessed: Option<ViewPair<'a, F>>,
127    /// The public values provided for constraint validation (e.g. inputs or outputs).
128    public_values: &'a [F],
129    /// A flag indicating whether this is the first row.
130    is_first_row: F,
131    /// A flag indicating whether this is the last row.
132    is_last_row: F,
133    /// A flag indicating whether this is a transition row (not the last row).
134    is_transition: F,
135    /// A view of the current and next permutation rows as a vertical pair.
136    permutation: ViewPair<'a, EF>,
137    /// The challenges used for the permutation argument.
138    permutation_challenges: &'a [EF],
139}
140
141impl<'a, F, EF> AirBuilder for DebugConstraintBuilderWithLookups<'a, F, EF>
142where
143    F: Field,
144    EF: ExtensionField<F>,
145{
146    type F = F;
147    type Expr = F;
148    type Var = F;
149    type M = ViewPair<'a, F>;
150
151    fn main(&self) -> Self::M {
152        self.main
153    }
154
155    fn is_first_row(&self) -> Self::Expr {
156        self.is_first_row
157    }
158
159    fn is_last_row(&self) -> Self::Expr {
160        self.is_last_row
161    }
162
163    /// # Panics
164    /// This function panics if `size` is not `2`.
165    fn is_transition_window(&self, size: usize) -> Self::Expr {
166        if size == 2 {
167            self.is_transition
168        } else {
169            panic!("only supports a window size of 2")
170        }
171    }
172
173    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
174        assert_eq!(
175            x.into(),
176            F::ZERO,
177            "constraints had nonzero value on row {}",
178            self.row_index
179        );
180    }
181
182    fn assert_eq<I1: Into<Self::Expr>, I2: Into<Self::Expr>>(&mut self, x: I1, y: I2) {
183        let x = x.into();
184        let y = y.into();
185        assert_eq!(
186            x, y,
187            "values didn't match on row {}: {} != {}",
188            self.row_index, x, y
189        );
190    }
191}
192
193impl<F: Field, EF: ExtensionField<F>> AirBuilderWithPublicValues
194    for DebugConstraintBuilderWithLookups<'_, F, EF>
195{
196    type PublicVar = Self::F;
197
198    fn public_values(&self) -> &[Self::F] {
199        self.public_values
200    }
201}
202
203impl<'a, F: Field, EF: ExtensionField<F>> ExtensionBuilder
204    for DebugConstraintBuilderWithLookups<'a, F, EF>
205{
206    type EF = EF;
207
208    type ExprEF = EF;
209
210    type VarEF = EF;
211
212    fn assert_zero_ext<I>(&mut self, x: I)
213    where
214        I: Into<Self::ExprEF>,
215    {
216        assert_eq!(
217            x.into(),
218            EF::ZERO,
219            "constraints had nonzero value on row {}",
220            self.row_index
221        );
222    }
223}
224
225impl<'a, F: Field, EF: ExtensionField<F>> PermutationAirBuilder
226    for DebugConstraintBuilderWithLookups<'a, F, EF>
227{
228    type MP = VerticalPair<RowMajorMatrixView<'a, EF>, RowMajorMatrixView<'a, EF>>;
229
230    type RandomVar = EF;
231
232    fn permutation(&self) -> Self::MP {
233        self.permutation
234    }
235
236    fn permutation_randomness(&self) -> &[Self::RandomVar] {
237        self.permutation_challenges
238    }
239}
240
241impl<'a, F: Field, EF: ExtensionField<F>> PairBuilder
242    for DebugConstraintBuilderWithLookups<'a, F, EF>
243{
244    fn preprocessed(&self) -> Self::M {
245        self.preprocessed
246            .unwrap_or_else(|| panic!("Missing preprocessed columns"))
247    }
248}
p3_batch_stark/check_constraints.rs

p3_batch_stark/
check_constraints.rs
