openvm_stark_backend/air_builders/symbolic/
mod.rs

1// Originally copied from uni-stark/src/symbolic_builder.rs to allow A: ?Sized
2
3use std::iter;
4
5use itertools::Itertools;
6use p3_air::{
7    Air, AirBuilder, AirBuilderWithPublicValues, BaseAirWithPublicValues, ExtensionBuilder,
8    PairBuilder,
9};
10use p3_field::Field;
11use p3_matrix::dense::RowMajorMatrix;
12use tracing::instrument;
13
14use self::{
15    symbolic_expression::SymbolicExpression,
16    symbolic_variable::{Entry, SymbolicVariable},
17};
18use super::PartitionedAirBuilder;
19use crate::{
20    interaction::{Interaction, InteractionBuilder, SymbolicInteraction},
21    keygen::types::TraceWidth,
22};
23
24mod dag;
25pub mod statistics;
26pub mod symbolic_expression;
27pub mod symbolic_variable;
28
29pub use dag::*;
30
31use crate::interaction::BusIndex;
32
33/// Symbolic constraints for a single AIR with interactions.
34#[derive(Clone, Debug)]
35pub struct SymbolicConstraints<F> {
36    /// All plain AIR constraints. These do **not** include interaction constraints that are proven
37    /// via LogUp-GKR.
38    pub constraints: Vec<SymbolicExpression<F>>,
39    /// Symbolic representation of interactions. These are converted into a LogUp fractional sum
40    /// which must be proven using GKR.
41    pub interactions: Vec<SymbolicInteraction<F>>,
42}
43
44impl<F: Field> SymbolicConstraints<F> {
45    pub fn max_constraint_degree(&self) -> usize {
46        iter::empty()
47            .chain(&self.constraints)
48            .chain(
49                self.interactions
50                    .iter()
51                    .flat_map(|i| iter::once(&i.count).chain(&i.message)),
52            )
53            .map(|expr| expr.degree_multiple())
54            .max()
55            .unwrap_or(0)
56    }
57
58    /// Returns the maximum field degree and count degree across all interactions
59    pub fn max_interaction_degrees(&self) -> (usize, usize) {
60        let max_field_degree = self
61            .interactions
62            .iter()
63            .map(|interaction| {
64                interaction
65                    .message
66                    .iter()
67                    .map(|field| field.degree_multiple())
68                    .max()
69                    .unwrap_or(0)
70            })
71            .max()
72            .unwrap_or(0);
73
74        let max_count_degree = self
75            .interactions
76            .iter()
77            .map(|interaction| interaction.count.degree_multiple())
78            .max()
79            .unwrap_or(0);
80
81        (max_field_degree, max_count_degree)
82    }
83}
84
85#[instrument(name = "evaluate constraints symbolically", skip_all, level = "debug")]
86pub fn get_symbolic_builder<F, R>(rap: &R, width: &TraceWidth) -> SymbolicRapBuilder<F>
87where
88    F: Field,
89    R: Air<SymbolicRapBuilder<F>> + BaseAirWithPublicValues<F> + ?Sized,
90{
91    let mut builder = SymbolicRapBuilder::new(width, rap.num_public_values());
92    Air::eval(rap, &mut builder);
93    builder
94}
95
96/// An `AirBuilder` for evaluating constraints symbolically, and recording them for later use.
97#[derive(Debug)]
98pub struct SymbolicRapBuilder<F> {
99    preprocessed: RowMajorMatrix<SymbolicVariable<F>>,
100    partitioned_main: Vec<RowMajorMatrix<SymbolicVariable<F>>>,
101    public_values: Vec<SymbolicVariable<F>>,
102    constraints: Vec<SymbolicExpression<F>>,
103    interactions: Vec<SymbolicInteraction<F>>,
104    trace_width: TraceWidth,
105}
106
107impl<F: Field> SymbolicRapBuilder<F> {
108    pub(crate) fn new(width: &TraceWidth, num_public_values: usize) -> Self {
109        let preprocessed_width = width.preprocessed.unwrap_or(0);
110        let prep_values = [0, 1]
111            .into_iter()
112            .flat_map(|offset| {
113                (0..width.preprocessed.unwrap_or(0))
114                    .map(move |index| SymbolicVariable::new(Entry::Preprocessed { offset }, index))
115            })
116            .collect();
117        let preprocessed = RowMajorMatrix::new(prep_values, preprocessed_width);
118
119        let mut partitioned_main: Vec<_> = width
120            .cached_mains
121            .iter()
122            .enumerate()
123            .map(|(part_index, &width)| gen_main_trace(part_index, width))
124            .collect();
125        if width.common_main != 0 {
126            partitioned_main.push(gen_main_trace(width.cached_mains.len(), width.common_main));
127        }
128
129        let public_values = (0..num_public_values)
130            .map(move |index| SymbolicVariable::new(Entry::Public, index))
131            .collect();
132
133        Self {
134            preprocessed,
135            partitioned_main,
136            public_values,
137            constraints: vec![],
138            interactions: vec![],
139            trace_width: width.clone(),
140        }
141    }
142
143    pub fn constraints(self) -> SymbolicConstraints<F> {
144        SymbolicConstraints {
145            constraints: self.constraints,
146            interactions: self.interactions,
147        }
148    }
149
150    pub fn num_public_values(&self) -> usize {
151        self.public_values.len()
152    }
153
154    pub fn width(&self) -> TraceWidth {
155        self.trace_width.clone()
156    }
157}
158
159impl<F: Field> AirBuilder for SymbolicRapBuilder<F> {
160    type F = F;
161    type Expr = SymbolicExpression<Self::F>;
162    type Var = SymbolicVariable<Self::F>;
163    type M = RowMajorMatrix<Self::Var>;
164
165    /// It is difficult to horizontally concatenate matrices when the main trace is partitioned, so
166    /// we disable this method in that case.
167    fn main(&self) -> Self::M {
168        if self.partitioned_main.len() == 1 {
169            self.partitioned_main[0].clone()
170        } else {
171            panic!("Main trace is either empty or partitioned. This function should not be used.")
172        }
173    }
174
175    fn is_first_row(&self) -> Self::Expr {
176        SymbolicExpression::IsFirstRow
177    }
178
179    fn is_last_row(&self) -> Self::Expr {
180        SymbolicExpression::IsLastRow
181    }
182
183    fn is_transition_window(&self, size: usize) -> Self::Expr {
184        if size == 2 {
185            SymbolicExpression::IsTransition
186        } else {
187            panic!("uni-stark only supports a window size of 2")
188        }
189    }
190
191    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
192        self.constraints.push(x.into());
193    }
194}
195
196impl<F: Field> PairBuilder for SymbolicRapBuilder<F> {
197    fn preprocessed(&self) -> Self::M {
198        self.preprocessed.clone()
199    }
200}
201
202impl<F: Field> ExtensionBuilder for SymbolicRapBuilder<F> {
203    type EF = F;
204    type ExprEF = SymbolicExpression<F>;
205    type VarEF = SymbolicVariable<F>;
206
207    fn assert_zero_ext<I>(&mut self, x: I)
208    where
209        I: Into<Self::ExprEF>,
210    {
211        self.constraints.push(x.into());
212    }
213}
214
215impl<F: Field> AirBuilderWithPublicValues for SymbolicRapBuilder<F> {
216    type PublicVar = SymbolicVariable<F>;
217
218    fn public_values(&self) -> &[Self::PublicVar] {
219        &self.public_values
220    }
221}
222
223impl<F: Field> InteractionBuilder for SymbolicRapBuilder<F> {
224    fn push_interaction<E: Into<Self::Expr>>(
225        &mut self,
226        bus_index: BusIndex,
227        fields: impl IntoIterator<Item = E>,
228        count: impl Into<Self::Expr>,
229        count_weight: u32,
230    ) {
231        let fields = fields.into_iter().map(|f| f.into()).collect();
232        let count = count.into();
233        self.interactions.push(Interaction {
234            bus_index,
235            message: fields,
236            count,
237            count_weight,
238        });
239    }
240
241    fn num_interactions(&self) -> usize {
242        self.interactions.len()
243    }
244
245    fn all_interactions(&self) -> &[Interaction<Self::Expr>] {
246        &self.interactions
247    }
248}
249
250impl<F: Field> PartitionedAirBuilder for SymbolicRapBuilder<F> {
251    fn cached_mains(&self) -> &[Self::M] {
252        &self.partitioned_main[..self.trace_width.cached_mains.len()]
253    }
254    fn common_main(&self) -> &Self::M {
255        assert_ne!(
256            self.trace_width.common_main, 0,
257            "AIR doesn't have a common main trace"
258        );
259        &self.partitioned_main[self.trace_width.cached_mains.len()]
260    }
261}
262
263#[allow(dead_code)]
264struct LocalOnlyChecker;
265
266#[allow(dead_code)]
267impl LocalOnlyChecker {
268    fn check_var<F: Field>(var: SymbolicVariable<F>) -> bool {
269        match var.entry {
270            Entry::Preprocessed { offset } => offset == 0,
271            Entry::Main { offset, .. } => offset == 0,
272            Entry::Public => true,
273            Entry::Challenge => true,
274        }
275    }
276
277    fn check_expr<F: Field>(expr: &SymbolicExpression<F>) -> bool {
278        match expr {
279            SymbolicExpression::Variable(var) => Self::check_var(*var),
280            SymbolicExpression::IsFirstRow => false,
281            SymbolicExpression::IsLastRow => false,
282            SymbolicExpression::IsTransition => false,
283            SymbolicExpression::Constant(_) => true,
284            SymbolicExpression::Add { x, y, .. } => Self::check_expr(x) && Self::check_expr(y),
285            SymbolicExpression::Sub { x, y, .. } => Self::check_expr(x) && Self::check_expr(y),
286            SymbolicExpression::Neg { x, .. } => Self::check_expr(x),
287            SymbolicExpression::Mul { x, y, .. } => Self::check_expr(x) && Self::check_expr(y),
288        }
289    }
290}
291
292fn gen_main_trace<F: Field>(
293    part_index: usize,
294    width: usize,
295) -> RowMajorMatrix<SymbolicVariable<F>> {
296    let mat_values = [0, 1]
297        .into_iter()
298        .flat_map(|offset| {
299            (0..width)
300                .map(move |index| SymbolicVariable::new(Entry::Main { part_index, offset }, index))
301        })
302        .collect_vec();
303    RowMajorMatrix::new(mat_values, width)
304}
openvm_stark_backend/air_builders/symbolic/mod.rs

openvm_stark_backend/air_builders/symbolic/
mod.rs
