openvm_stark_backend/air_builders/debug/
mod.rs

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
33/// phase 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
108    /// we disable this method in that case.
109    fn main(&self) -> Self::M {
110        if self.partitioned_main.len() == 1 {
111            self.partitioned_main[0]
112        } else {
113            panic!("Main trace is either empty or partitioned. This function should not be used.")
114        }
115    }
116
117    fn is_first_row(&self) -> Self::Expr {
118        self.is_first_row
119    }
120
121    fn is_last_row(&self) -> Self::Expr {
122        self.is_last_row
123    }
124
125    fn is_transition_window(&self, size: usize) -> Self::Expr {
126        if size == 2 {
127            self.is_transition
128        } else {
129            panic!("only supports a window size of 2")
130        }
131    }
132
133    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
134        assert_eq!(
135            x.into(),
136            Val::<SC>::ZERO,
137            "constraints had nonzero value on air {},row {}",
138            self.air_name,
139            self.row_index
140        );
141    }
142
143    fn assert_eq<I1: Into<Self::Expr>, I2: Into<Self::Expr>>(&mut self, x: I1, y: I2) {
144        let x = x.into();
145        let y = y.into();
146        assert_eq!(
147            x, y,
148            "values didn't match on air {}, row {}: {} != {}",
149            self.air_name, self.row_index, x, y
150        );
151    }
152}
153
154impl<SC> PairBuilder for DebugConstraintBuilder<'_, SC>
155where
156    SC: StarkGenericConfig,
157{
158    fn preprocessed(&self) -> Self::M {
159        self.preprocessed
160    }
161}
162
163impl<SC> ExtensionBuilder for DebugConstraintBuilder<'_, SC>
164where
165    SC: StarkGenericConfig,
166{
167    type EF = SC::Challenge;
168    type ExprEF = SC::Challenge;
169    type VarEF = SC::Challenge;
170
171    fn assert_zero_ext<I>(&mut self, x: I)
172    where
173        I: Into<Self::ExprEF>,
174    {
175        assert_eq!(
176            x.into(),
177            SC::Challenge::ZERO,
178            "constraints had nonzero value on row {}",
179            self.row_index
180        );
181    }
182
183    fn assert_eq_ext<I1, I2>(&mut self, x: I1, y: I2)
184    where
185        I1: Into<Self::ExprEF>,
186        I2: Into<Self::ExprEF>,
187    {
188        let x = x.into();
189        let y = y.into();
190        assert_eq!(
191            x, y,
192            "values didn't match on air {}, row {}: {} != {}",
193            self.air_name, self.row_index, x, y
194        );
195    }
196}
197
198impl<'a, SC> PermutationAirBuilder for DebugConstraintBuilder<'a, SC>
199where
200    SC: StarkGenericConfig,
201{
202    type MP = ViewPair<'a, SC::Challenge>;
203
204    type RandomVar = SC::Challenge;
205
206    fn permutation(&self) -> Self::MP {
207        *self
208            .after_challenge
209            .first()
210            .expect("Challenge phase not supported")
211    }
212
213    fn permutation_randomness(&self) -> &[Self::EF] {
214        self.challenges
215            .first()
216            .expect("Challenge phase not supported")
217    }
218}
219
220impl<SC> AirBuilderWithPublicValues for DebugConstraintBuilder<'_, SC>
221where
222    SC: StarkGenericConfig,
223{
224    type PublicVar = Val<SC>;
225
226    fn public_values(&self) -> &[Self::F] {
227        self.public_values
228    }
229}
230
231impl<SC> PermutationAirBuilderWithExposedValues for DebugConstraintBuilder<'_, SC>
232where
233    SC: StarkGenericConfig,
234{
235    fn permutation_exposed_values(&self) -> &[Self::EF] {
236        self.exposed_values_after_challenge
237            .first()
238            .expect("Challenge phase not supported")
239    }
240}
241
242impl<SC> PartitionedAirBuilder for DebugConstraintBuilder<'_, SC>
243where
244    SC: StarkGenericConfig,
245{
246    fn cached_mains(&self) -> &[Self::M] {
247        let mut num_cached_mains = self.partitioned_main.len();
248        if self.has_common_main {
249            num_cached_mains -= 1;
250        }
251        &self.partitioned_main[..num_cached_mains]
252    }
253    fn common_main(&self) -> &Self::M {
254        assert!(self.has_common_main, "AIR doesn't have a common main trace");
255        self.partitioned_main.last().unwrap()
256    }
257}
258
259// No-op implementation
260impl<SC> InteractionBuilder for DebugConstraintBuilder<'_, SC>
261where
262    SC: StarkGenericConfig,
263{
264    fn push_interaction<E: Into<Self::Expr>>(
265        &mut self,
266        _bus_index: BusIndex,
267        _fields: impl IntoIterator<Item = E>,
268        _count: impl Into<Self::Expr>,
269        _count_weight: u32,
270    ) {
271        // no-op, interactions are debugged elsewhere
272    }
273
274    fn num_interactions(&self) -> usize {
275        0
276    }
277
278    fn all_interactions(&self) -> &[Interaction<Self::Expr>] {
279        &[]
280    }
281}
282
283// No-op
284impl<SC: StarkGenericConfig> InteractionPhaseAirBuilder for DebugConstraintBuilder<'_, SC> {
285    fn finalize_interactions(&mut self) {}
286
287    fn max_constraint_degree(&self) -> usize {
288        0
289    }
290
291    fn rap_phase_seq_kind(&self) -> RapPhaseSeqKind {
292        self.rap_phase_seq_kind
293    }
294
295    fn symbolic_interactions(&self) -> Vec<SymbolicInteraction<Val<SC>>> {
296        vec![]
297    }
298}
openvm_stark_backend/air_builders/debug/mod.rs

openvm_stark_backend/air_builders/debug/
mod.rs
