openvm_stark_backend/air_builders/debug/
mod.rs

Help
1use std::sync::{Arc, Mutex};
2
3use itertools::{izip, Itertools};
4use p3_air::{
5    AirBuilder, AirBuilderWithPublicValues, ExtensionBuilder, PairBuilder, PermutationAirBuilder,
6};
7use p3_field::FieldAlgebra;
8use p3_matrix::{dense::RowMajorMatrixView, stack::VerticalPair};
9
10use super::{symbolic::SymbolicConstraints, PartitionedAirBuilder, ViewPair};
11use crate::{
12    config::{StarkGenericConfig, Val},
13    interaction::{
14        rap::InteractionPhaseAirBuilder, Interaction, InteractionBuilder, RapPhaseSeqKind,
15        SymbolicInteraction,
16    },
17    keygen::types::StarkProvingKey,
18    rap::{AnyRap, PermutationAirBuilderWithExposedValues},
19};
20
21mod check_constraints;
22
23use check_constraints::*;
24
25use crate::interaction::BusIndex;
26
27thread_local! {
28   pub static USE_DEBUG_BUILDER: Arc<Mutex<bool>> = Arc::new(Mutex::new(true));
29}
30
31/// The debugging will check the main AIR constraints and then separately check LogUp constraints by
32/// checking the actual multiset equalities. Currently it will not debug check any after challenge phase
33/// constraints for implementation simplicity.
34#[allow(dead_code)]
35#[allow(clippy::too_many_arguments)]
36pub fn debug_constraints_and_interactions<SC: StarkGenericConfig>(
37    airs: &[Arc<dyn AnyRap<SC>>],
38    pk: &[StarkProvingKey<SC>],
39    main_views_per_air: &[Vec<RowMajorMatrixView<'_, Val<SC>>>],
40    public_values_per_air: &[Vec<Val<SC>>],
41) {
42    USE_DEBUG_BUILDER.with(|debug| {
43        if *debug.lock().unwrap() {
44            let preprocessed = izip!(airs, pk, main_views_per_air, public_values_per_air,)
45                .map(|(rap, pk, main, public_values)| {
46                    let preprocessed_trace = pk
47                        .preprocessed_data
48                        .as_ref()
49                        .map(|data| data.trace.as_view());
50                    tracing::debug!("Checking constraints for {}", rap.name());
51                    check_constraints(
52                        rap.as_ref(),
53                        &rap.name(),
54                        &preprocessed_trace,
55                        main,
56                        public_values,
57                    );
58                    preprocessed_trace
59                })
60                .collect_vec();
61
62            let (air_names, interactions): (Vec<_>, Vec<_>) = pk
63                .iter()
64                .map(|pk| {
65                    let sym_constraints = SymbolicConstraints::from(&pk.vk.symbolic_constraints);
66                    (pk.air_name.clone(), sym_constraints.interactions)
67                })
68                .unzip();
69            check_logup(
70                &air_names,
71                &interactions,
72                &preprocessed,
73                main_views_per_air,
74                public_values_per_air,
75            );
76        }
77    });
78}
79
80/// An `AirBuilder` which asserts that each constraint is zero, allowing any failed constraints to
81/// be detected early.
82pub struct DebugConstraintBuilder<'a, SC: StarkGenericConfig> {
83    pub air_name: &'a str,
84    pub row_index: usize,
85    pub preprocessed: ViewPair<'a, Val<SC>>,
86    pub partitioned_main: Vec<ViewPair<'a, Val<SC>>>,
87    pub after_challenge: Vec<ViewPair<'a, SC::Challenge>>,
88    pub challenges: &'a [Vec<SC::Challenge>],
89    pub is_first_row: Val<SC>,
90    pub is_last_row: Val<SC>,
91    pub is_transition: Val<SC>,
92    pub public_values: &'a [Val<SC>],
93    pub exposed_values_after_challenge: &'a [Vec<SC::Challenge>],
94    pub rap_phase_seq_kind: RapPhaseSeqKind,
95    pub has_common_main: bool,
96}
97
98impl<'a, SC> AirBuilder for DebugConstraintBuilder<'a, SC>
99where
100    SC: StarkGenericConfig,
101{
102    type F = Val<SC>;
103    type Expr = Val<SC>;
104    type Var = Val<SC>;
105    type M = VerticalPair<RowMajorMatrixView<'a, Val<SC>>, RowMajorMatrixView<'a, Val<SC>>>;
106
107    /// It is difficult to horizontally concatenate matrices when the main trace is partitioned, so we disable this method in that case.
108    fn main(&self) -> Self::M {
109        if self.partitioned_main.len() == 1 {
110            self.partitioned_main[0]
111        } else {
112            panic!("Main trace is either empty or partitioned. This function should not be used.")
113        }
114    }
115
116    fn is_first_row(&self) -> Self::Expr {
117        self.is_first_row
118    }
119
120    fn is_last_row(&self) -> Self::Expr {
121        self.is_last_row
122    }
123
124    fn is_transition_window(&self, size: usize) -> Self::Expr {
125        if size == 2 {
126            self.is_transition
127        } else {
128            panic!("only supports a window size of 2")
129        }
130    }
131
132    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
133        assert_eq!(
134            x.into(),
135            Val::<SC>::ZERO,
136            "constraints had nonzero value on air {},row {}",
137            self.air_name,
138            self.row_index
139        );
140    }
141
142    fn assert_eq<I1: Into<Self::Expr>, I2: Into<Self::Expr>>(&mut self, x: I1, y: I2) {
143        let x = x.into();
144        let y = y.into();
145        assert_eq!(
146            x, y,
147            "values didn't match on air {}, row {}: {} != {}",
148            self.air_name, self.row_index, x, y
149        );
150    }
151}
152
153impl<SC> PairBuilder for DebugConstraintBuilder<'_, SC>
154where
155    SC: StarkGenericConfig,
156{
157    fn preprocessed(&self) -> Self::M {
158        self.preprocessed
159    }
160}
161
162impl<SC> ExtensionBuilder for DebugConstraintBuilder<'_, SC>
163where
164    SC: StarkGenericConfig,
165{
166    type EF = SC::Challenge;
167    type ExprEF = SC::Challenge;
168    type VarEF = SC::Challenge;
169
170    fn assert_zero_ext<I>(&mut self, x: I)
171    where
172        I: Into<Self::ExprEF>,
173    {
174        assert_eq!(
175            x.into(),
176            SC::Challenge::ZERO,
177            "constraints had nonzero value on row {}",
178            self.row_index
179        );
180    }
181
182    fn assert_eq_ext<I1, I2>(&mut self, x: I1, y: I2)
183    where
184        I1: Into<Self::ExprEF>,
185        I2: Into<Self::ExprEF>,
186    {
187        let x = x.into();
188        let y = y.into();
189        assert_eq!(
190            x, y,
191            "values didn't match on air {}, row {}: {} != {}",
192            self.air_name, self.row_index, x, y
193        );
194    }
195}
196
197impl<'a, SC> PermutationAirBuilder for DebugConstraintBuilder<'a, SC>
198where
199    SC: StarkGenericConfig,
200{
201    type MP = ViewPair<'a, SC::Challenge>;
202
203    type RandomVar = SC::Challenge;
204
205    fn permutation(&self) -> Self::MP {
206        *self
207            .after_challenge
208            .first()
209            .expect("Challenge phase not supported")
210    }
211
212    fn permutation_randomness(&self) -> &[Self::EF] {
213        self.challenges
214            .first()
215            .expect("Challenge phase not supported")
216    }
217}
218
219impl<SC> AirBuilderWithPublicValues for DebugConstraintBuilder<'_, SC>
220where
221    SC: StarkGenericConfig,
222{
223    type PublicVar = Val<SC>;
224
225    fn public_values(&self) -> &[Self::F] {
226        self.public_values
227    }
228}
229
230impl<SC> PermutationAirBuilderWithExposedValues for DebugConstraintBuilder<'_, SC>
231where
232    SC: StarkGenericConfig,
233{
234    fn permutation_exposed_values(&self) -> &[Self::EF] {
235        self.exposed_values_after_challenge
236            .first()
237            .expect("Challenge phase not supported")
238    }
239}
240
241impl<SC> PartitionedAirBuilder for DebugConstraintBuilder<'_, SC>
242where
243    SC: StarkGenericConfig,
244{
245    fn cached_mains(&self) -> &[Self::M] {
246        let mut num_cached_mains = self.partitioned_main.len();
247        if self.has_common_main {
248            num_cached_mains -= 1;
249        }
250        &self.partitioned_main[..num_cached_mains]
251    }
252    fn common_main(&self) -> &Self::M {
253        assert!(self.has_common_main, "AIR doesn't have a common main trace");
254        self.partitioned_main.last().unwrap()
255    }
256}
257
258// No-op implementation
259impl<SC> InteractionBuilder for DebugConstraintBuilder<'_, SC>
260where
261    SC: StarkGenericConfig,
262{
263    fn push_interaction<E: Into<Self::Expr>>(
264        &mut self,
265        _bus_index: BusIndex,
266        _fields: impl IntoIterator<Item = E>,
267        _count: impl Into<Self::Expr>,
268        _count_weight: u32,
269    ) {
270        // no-op, interactions are debugged elsewhere
271    }
272
273    fn num_interactions(&self) -> usize {
274        0
275    }
276
277    fn all_interactions(&self) -> &[Interaction<Self::Expr>] {
278        &[]
279    }
280}
281
282// No-op
283impl<SC: StarkGenericConfig> InteractionPhaseAirBuilder for DebugConstraintBuilder<'_, SC> {
284    fn finalize_interactions(&mut self) {}
285
286    fn max_constraint_degree(&self) -> usize {
287        0
288    }
289
290    fn rap_phase_seq_kind(&self) -> RapPhaseSeqKind {
291        self.rap_phase_seq_kind
292    }
293
294    fn symbolic_interactions(&self) -> Vec<SymbolicInteraction<Val<SC>>> {
295        vec![]
296    }
297}
openvm_stark_backend/air_builders/debug/mod.rs

openvm_stark_backend/air_builders/debug/
mod.rs
