openvm_stark_backend/air_builders/debug/
check_constraints.rs

1use std::sync::Arc;
2
3use itertools::{izip, Itertools};
4use p3_air::{Air, BaseAir};
5use p3_field::{Field, PrimeCharacteristicRing};
6use p3_matrix::{
7    dense::{RowMajorMatrix, RowMajorMatrixView},
8    stack::VerticalPair,
9    Matrix,
10};
11use p3_maybe_rayon::prelude::*;
12
13use crate::{
14    air_builders::{
15        debug::{DebugConstraintBuilder, USE_DEBUG_BUILDER},
16        symbolic::SymbolicConstraints,
17    },
18    config::{StarkProtocolConfig, Val},
19    interaction::{
20        debug::{generate_logical_interactions, LogicalInteractions},
21        SymbolicInteraction,
22    },
23    keygen::types::StarkProvingKey,
24    AirRef, PartitionedBaseAir,
25};
26
27/// Raw input data for debugging a single AIR.
28pub struct AirProofRawInput<F> {
29    pub cached_mains: Vec<Arc<RowMajorMatrix<F>>>,
30    pub common_main: Option<Arc<RowMajorMatrix<F>>>,
31    pub public_values: Vec<F>,
32}
33
34/// Check that all constraints vanish on the subgroup.
35#[allow(clippy::too_many_arguments)]
36pub fn check_constraints<R, SC>(
37    rap: &R,
38    rap_name: &str,
39    preprocessed: &Option<RowMajorMatrixView<Val<SC>>>,
40    partitioned_main: &[RowMajorMatrixView<Val<SC>>],
41    public_values: &[Val<SC>],
42) where
43    R: for<'a> Air<DebugConstraintBuilder<'a, SC>>
44        + BaseAir<Val<SC>>
45        + PartitionedBaseAir<Val<SC>>
46        + ?Sized,
47    SC: StarkProtocolConfig,
48{
49    let height = partitioned_main[0].height();
50    assert!(partitioned_main.iter().all(|mat| mat.height() == height));
51
52    // Check that constraints are satisfied.
53    (0..height).into_par_iter().for_each(|i| {
54        let i_next = (i + 1) % height;
55
56        let (preprocessed_local, preprocessed_next) = preprocessed
57            .as_ref()
58            .map(|preprocessed| {
59                (
60                    preprocessed.row_slice(i).unwrap().to_vec(),
61                    preprocessed.row_slice(i_next).unwrap().to_vec(),
62                )
63            })
64            .unwrap_or((vec![], vec![]));
65
66        let partitioned_main_row_pair = partitioned_main
67            .iter()
68            .map(|part| (part.row_slice(i).unwrap(), part.row_slice(i_next).unwrap()))
69            .collect::<Vec<_>>();
70        let partitioned_main = partitioned_main_row_pair
71            .iter()
72            .map(|(local, next)| {
73                VerticalPair::new(
74                    RowMajorMatrixView::new_row(local),
75                    RowMajorMatrixView::new_row(next),
76                )
77            })
78            .collect::<Vec<_>>();
79
80        let mut builder = DebugConstraintBuilder {
81            air_name: rap_name,
82            row_index: i,
83            preprocessed: VerticalPair::new(
84                RowMajorMatrixView::new_row(preprocessed_local.as_slice()),
85                RowMajorMatrixView::new_row(preprocessed_next.as_slice()),
86            ),
87            partitioned_main,
88            public_values,
89            is_first_row: Val::<SC>::ZERO,
90            is_last_row: Val::<SC>::ZERO,
91            is_transition: Val::<SC>::ONE,
92            has_common_main: rap.common_main_width() > 0,
93        };
94        if i == 0 {
95            builder.is_first_row = Val::<SC>::ONE;
96        }
97        if i == height - 1 {
98            builder.is_last_row = Val::<SC>::ONE;
99            builder.is_transition = Val::<SC>::ZERO;
100        }
101
102        rap.eval(&mut builder);
103    });
104}
105
106pub fn check_logup<F: Field>(
107    air_names: &[String],
108    interactions: &[Vec<SymbolicInteraction<F>>],
109    preprocessed: &[Option<RowMajorMatrixView<F>>],
110    partitioned_main: &[Vec<RowMajorMatrixView<F>>],
111    public_values: &[Vec<F>],
112) {
113    let mut logical_interactions = LogicalInteractions::<F>::default();
114    for (air_idx, (interactions, preprocessed, partitioned_main, public_values)) in
115        izip!(interactions, preprocessed, partitioned_main, public_values).enumerate()
116    {
117        generate_logical_interactions(
118            air_idx,
119            interactions,
120            preprocessed,
121            partitioned_main,
122            public_values,
123            &mut logical_interactions,
124        );
125    }
126
127    let mut logup_failed = false;
128    // For each bus, check each `fields` key by summing up multiplicities.
129    for (bus_idx, bus_interactions) in logical_interactions.at_bus.into_iter() {
130        for (fields, connections) in bus_interactions.into_iter() {
131            let sum: F = connections.iter().map(|(_, count)| *count).sum();
132            if !sum.is_zero() {
133                logup_failed = true;
134                println!(
135                    "Bus {} failed to balance the multiplicities for fields={:?}. The bus connections for this were:",
136                    bus_idx, fields
137                );
138                for (air_idx, count) in connections {
139                    println!(
140                        "   Air idx: {}, Air name: {}, count: {:?}",
141                        air_idx, air_names[air_idx], count
142                    );
143                }
144            }
145        }
146    }
147    if logup_failed {
148        panic!("LogUp multiset equality check failed.");
149    }
150}
151
152/// The debugging will check the main AIR constraints and then separately check LogUp constraints by
153/// checking the actual multiset equalities. Currently it will not debug check any after challenge
154/// phase constraints for implementation simplicity.
155#[allow(clippy::too_many_arguments)]
156pub fn debug_constraints_and_interactions<SC: StarkProtocolConfig>(
157    airs: &[AirRef<SC>],
158    pk: &[&StarkProvingKey<SC>],
159    inputs: &[AirProofRawInput<SC::F>],
160) {
161    USE_DEBUG_BUILDER.with(|debug| {
162        if *debug.lock().unwrap() {
163            let (main_parts_per_air, pvs_per_air): (Vec<_>, Vec<_>) = inputs
164                .iter()
165                .map(|input| {
166                    let mut main_parts = input
167                        .cached_mains
168                        .iter()
169                        .map(|trace| trace.as_view())
170                        .collect_vec();
171                    if let Some(trace) = input.common_main.as_ref() {
172                        main_parts.push(trace.as_view());
173                    }
174                    (main_parts, input.public_values.clone())
175                })
176                .unzip();
177            let preprocessed = izip!(airs, pk, &main_parts_per_air, &pvs_per_air)
178                .map(|(air, pk, main_parts, pvs)| {
179                    let preprocessed_trace = pk
180                        .preprocessed_data
181                        .as_ref()
182                        .map(|data| data.mat_view(0).to_row_major_matrix());
183                    tracing::debug!("Checking constraints for {}", air.name());
184                    check_constraints(
185                        air.as_ref(),
186                        &air.name(),
187                        &preprocessed_trace.as_ref().map(|t| t.as_view()),
188                        main_parts,
189                        pvs,
190                    );
191                    preprocessed_trace
192                })
193                .collect_vec();
194
195            let (air_names, interactions): (Vec<_>, Vec<_>) = pk
196                .iter()
197                .map(|pk| {
198                    let sym_constraints = SymbolicConstraints::from(&pk.vk.symbolic_constraints);
199                    (pk.air_name.clone(), sym_constraints.interactions)
200                })
201                .unzip();
202            let preprocessed_views = preprocessed
203                .iter()
204                .map(|t| t.as_ref().map(|t| t.as_view()))
205                .collect_vec();
206            check_logup(
207                &air_names,
208                &interactions,
209                &preprocessed_views,
210                &main_parts_per_air,
211                &pvs_per_air,
212            );
213        }
214    });
215}
openvm_stark_backend/air_builders/debug/check_constraints.rs

openvm_stark_backend/air_builders/debug/
check_constraints.rs
