p3_uni_stark/
symbolic_builder.rs

Help
1use alloc::vec;
2use alloc::vec::Vec;
3
4use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, PairBuilder};
5use p3_field::Field;
6use p3_matrix::dense::RowMajorMatrix;
7use p3_util::log2_ceil_usize;
8use tracing::instrument;
9
10use crate::symbolic_expression::SymbolicExpression;
11use crate::symbolic_variable::SymbolicVariable;
12use crate::Entry;
13
14#[instrument(name = "infer log of constraint degree", skip_all)]
15pub fn get_log_quotient_degree<F, A>(
16    air: &A,
17    preprocessed_width: usize,
18    num_public_values: usize,
19) -> usize
20where
21    F: Field,
22    A: Air<SymbolicAirBuilder<F>>,
23{
24    // We pad to at least degree 2, since a quotient argument doesn't make sense with smaller degrees.
25    let constraint_degree =
26        get_max_constraint_degree(air, preprocessed_width, num_public_values).max(2);
27
28    // The quotient's actual degree is approximately (max_constraint_degree - 1) n,
29    // where subtracting 1 comes from division by the zerofier.
30    // But we pad it to a power of two so that we can efficiently decompose the quotient.
31    log2_ceil_usize(constraint_degree - 1)
32}
33
34#[instrument(name = "infer constraint degree", skip_all, level = "debug")]
35pub fn get_max_constraint_degree<F, A>(
36    air: &A,
37    preprocessed_width: usize,
38    num_public_values: usize,
39) -> usize
40where
41    F: Field,
42    A: Air<SymbolicAirBuilder<F>>,
43{
44    get_symbolic_constraints(air, preprocessed_width, num_public_values)
45        .iter()
46        .map(|c| c.degree_multiple())
47        .max()
48        .unwrap_or(0)
49}
50
51#[instrument(name = "evaluate constraints symbolically", skip_all, level = "debug")]
52pub fn get_symbolic_constraints<F, A>(
53    air: &A,
54    preprocessed_width: usize,
55    num_public_values: usize,
56) -> Vec<SymbolicExpression<F>>
57where
58    F: Field,
59    A: Air<SymbolicAirBuilder<F>>,
60{
61    let mut builder = SymbolicAirBuilder::new(preprocessed_width, air.width(), num_public_values);
62    air.eval(&mut builder);
63    builder.constraints()
64}
65
66/// An `AirBuilder` for evaluating constraints symbolically, and recording them for later use.
67#[derive(Debug)]
68pub struct SymbolicAirBuilder<F: Field> {
69    preprocessed: RowMajorMatrix<SymbolicVariable<F>>,
70    main: RowMajorMatrix<SymbolicVariable<F>>,
71    public_values: Vec<SymbolicVariable<F>>,
72    constraints: Vec<SymbolicExpression<F>>,
73}
74
75impl<F: Field> SymbolicAirBuilder<F> {
76    pub(crate) fn new(preprocessed_width: usize, width: usize, num_public_values: usize) -> Self {
77        let prep_values = [0, 1]
78            .into_iter()
79            .flat_map(|offset| {
80                (0..preprocessed_width)
81                    .map(move |index| SymbolicVariable::new(Entry::Preprocessed { offset }, index))
82            })
83            .collect();
84        let main_values = [0, 1]
85            .into_iter()
86            .flat_map(|offset| {
87                (0..width).map(move |index| SymbolicVariable::new(Entry::Main { offset }, index))
88            })
89            .collect();
90        let public_values = (0..num_public_values)
91            .map(move |index| SymbolicVariable::new(Entry::Public, index))
92            .collect();
93        Self {
94            preprocessed: RowMajorMatrix::new(prep_values, preprocessed_width),
95            main: RowMajorMatrix::new(main_values, width),
96            public_values,
97            constraints: vec![],
98        }
99    }
100
101    pub(crate) fn constraints(self) -> Vec<SymbolicExpression<F>> {
102        self.constraints
103    }
104}
105
106impl<F: Field> AirBuilder for SymbolicAirBuilder<F> {
107    type F = F;
108    type Expr = SymbolicExpression<F>;
109    type Var = SymbolicVariable<F>;
110    type M = RowMajorMatrix<Self::Var>;
111
112    fn main(&self) -> Self::M {
113        self.main.clone()
114    }
115
116    fn is_first_row(&self) -> Self::Expr {
117        SymbolicExpression::IsFirstRow
118    }
119
120    fn is_last_row(&self) -> Self::Expr {
121        SymbolicExpression::IsLastRow
122    }
123
124    fn is_transition_window(&self, size: usize) -> Self::Expr {
125        if size == 2 {
126            SymbolicExpression::IsTransition
127        } else {
128            panic!("uni-stark only supports a window size of 2")
129        }
130    }
131
132    fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
133        self.constraints.push(x.into());
134    }
135}
136
137impl<F: Field> AirBuilderWithPublicValues for SymbolicAirBuilder<F> {
138    type PublicVar = SymbolicVariable<F>;
139    fn public_values(&self) -> &[Self::PublicVar] {
140        &self.public_values
141    }
142}
143
144impl<F: Field> PairBuilder for SymbolicAirBuilder<F> {
145    fn preprocessed(&self) -> Self::M {
146        self.preprocessed.clone()
147    }
148}
p3_uni_stark/symbolic_builder.rs

p3_uni_stark/
symbolic_builder.rs
