p3_air/symbolic/
mod.rs

1//! Symbolic expression types for AIR constraint representation.
2
3mod builder;
4mod expression;
5pub(crate) mod expression_ext;
6mod flatten;
7mod variable;
8
9use alloc::collections::BTreeMap;
10use alloc::sync::Arc;
11use core::iter::{Product, Sum};
12use core::ops;
13
14pub use builder::*;
15pub use expression::{BaseLeaf, SymbolicExpression};
16pub use expression_ext::{ExtLeaf, SymbolicExpressionExt};
17use p3_field::{Dup, ExtensionField, Field, PrimeCharacteristicRing};
18pub use variable::{BaseEntry, ExtEntry, SymbolicVariable, SymbolicVariableExt};
19
20/// Properties that leaf nodes must provide for the generic expression tree.
21///
22/// Both [`BaseLeaf`] (base-field) and
23/// [`ExtLeaf`] (extension-field) implement this trait,
24/// enabling [`SymbolicExpr`] to handle constant folding, degree tracking, and
25/// arithmetic generically.
26pub trait SymLeaf: Clone + core::fmt::Debug {
27    /// The base field type used for constant folding.
28    type F: Field;
29
30    const ZERO: Self;
31    const ONE: Self;
32    const TWO: Self;
33    const NEG_ONE: Self;
34
35    /// Returns the degree multiple of this leaf.
36    fn degree_multiple(&self) -> usize;
37
38    /// Returns the exact polynomial degree of this leaf over a trace of length
39    /// `trace_len`, given the period of each periodic column.
40    fn poly_degree(&self, trace_len: usize, periodic_periods: &[usize]) -> usize;
41
42    /// Try to view this leaf as a base-field constant.
43    fn as_const(&self) -> Option<&Self::F>;
44
45    /// Create a leaf from a base-field constant.
46    fn from_const(c: Self::F) -> Self;
47}
48
49/// A symbolic expression tree, generic over its leaf type `A`.
50///
51/// This enum captures the shared tree structure — Add/Sub/Neg/Mul nodes with
52/// `Arc`-wrapped children and cached degree multiples — used by both base-field
53/// and extension-field symbolic expressions.
54///
55/// Concrete types are provided via type aliases:
56/// - [`SymbolicExpression<F>`] = `SymbolicExpr<BaseLeaf<F>>` (base-field constraints)
57/// - [`SymbolicExpressionExt<F, EF>`] = `SymbolicExpr<ExtLeaf<F, EF>>` (extension-field constraints)
58#[derive(Clone, Debug)]
59pub enum SymbolicExpr<A> {
60    /// A leaf node (variable, constant, selector, or lifted sub-expression).
61    Leaf(A),
62
63    /// Addition of two sub-expressions.
64    Add {
65        x: Arc<Self>,
66        y: Arc<Self>,
67        degree_multiple: usize,
68    },
69
70    /// Subtraction of two sub-expressions.
71    Sub {
72        x: Arc<Self>,
73        y: Arc<Self>,
74        degree_multiple: usize,
75    },
76
77    /// Negation of a sub-expression.
78    Neg {
79        x: Arc<Self>,
80        degree_multiple: usize,
81    },
82
83    /// Multiplication of two sub-expressions.
84    Mul {
85        x: Arc<Self>,
86        y: Arc<Self>,
87        degree_multiple: usize,
88    },
89}
90
91impl<A: SymLeaf> SymbolicExpr<A> {
92    /// Returns the degree multiple of this expression.
93    pub fn degree_multiple(&self) -> usize {
94        match self {
95            Self::Leaf(a) => a.degree_multiple(),
96            Self::Add {
97                degree_multiple, ..
98            }
99            | Self::Sub {
100                degree_multiple, ..
101            }
102            | Self::Neg {
103                degree_multiple, ..
104            }
105            | Self::Mul {
106                degree_multiple, ..
107            } => *degree_multiple,
108        }
109    }
110
111    /// Returns the exact polynomial degree of this expression over a trace of
112    /// length `trace_len`, given the period of each periodic column (indexed by
113    /// periodic column index).
114    ///
115    /// This is trace-size aware: it treats the transition selector as the linear
116    /// polynomial it is and accounts for the reduced degree of periodic columns,
117    /// unlike the trace-size-independent [`Self::degree_multiple`].
118    pub fn poly_degree(&self, trace_len: usize, periodic_periods: &[usize]) -> usize {
119        // The expression is a DAG: arithmetic nodes share `Arc` children, so a naive
120        // recursion would revisit shared subtrees exponentially. Memoize on node
121        // identity to keep this linear in the number of distinct nodes.
122        let mut cache: BTreeMap<*const Self, usize> = BTreeMap::new();
123        self.poly_degree_memo(trace_len, periodic_periods, &mut cache)
124    }
125
126    fn poly_degree_memo(
127        &self,
128        trace_len: usize,
129        periodic_periods: &[usize],
130        cache: &mut BTreeMap<*const Self, usize>,
131    ) -> usize {
132        match self {
133            Self::Leaf(a) => a.poly_degree(trace_len, periodic_periods),
134            Self::Add { x, y, .. } | Self::Sub { x, y, .. } => {
135                Self::child_poly_degree(x, trace_len, periodic_periods, cache).max(
136                    Self::child_poly_degree(y, trace_len, periodic_periods, cache),
137                )
138            }
139            Self::Neg { x, .. } => Self::child_poly_degree(x, trace_len, periodic_periods, cache),
140            Self::Mul { x, y, .. } => {
141                Self::child_poly_degree(x, trace_len, periodic_periods, cache)
142                    + Self::child_poly_degree(y, trace_len, periodic_periods, cache)
143            }
144        }
145    }
146
147    /// Degree of an `Arc`-shared child, looked up by pointer identity so each
148    /// distinct node is evaluated at most once.
149    fn child_poly_degree(
150        node: &Arc<Self>,
151        trace_len: usize,
152        periodic_periods: &[usize],
153        cache: &mut BTreeMap<*const Self, usize>,
154    ) -> usize {
155        let key = Arc::as_ptr(node);
156        if let Some(&degree) = cache.get(&key) {
157            return degree;
158        }
159        let degree = node.poly_degree_memo(trace_len, periodic_periods, cache);
160        cache.insert(key, degree);
161        degree
162    }
163
164    /// Try to view this expression as a base-field constant.
165    fn as_const(&self) -> Option<&A::F> {
166        match self {
167            Self::Leaf(a) => a.as_const(),
168            _ => None,
169        }
170    }
171
172    /// Addition with constant folding and zero-identity elimination.
173    fn sym_add(self, rhs: Self) -> Self {
174        if let (Some(&a), Some(&b)) = (self.as_const(), rhs.as_const()) {
175            return Self::Leaf(A::from_const(a + b));
176        }
177        if self.as_const().is_some_and(|c| c.is_zero()) {
178            return rhs;
179        }
180        if rhs.as_const().is_some_and(|c| c.is_zero()) {
181            return self;
182        }
183        let dm = self.degree_multiple().max(rhs.degree_multiple());
184        Self::Add {
185            x: Arc::new(self),
186            y: Arc::new(rhs),
187            degree_multiple: dm,
188        }
189    }
190
191    /// Subtraction with constant folding and zero-identity elimination.
192    fn sym_sub(self, rhs: Self) -> Self {
193        if let (Some(&a), Some(&b)) = (self.as_const(), rhs.as_const()) {
194            return Self::Leaf(A::from_const(a - b));
195        }
196        if self.as_const().is_some_and(|c| c.is_zero()) {
197            return rhs.sym_neg();
198        }
199        if rhs.as_const().is_some_and(|c| c.is_zero()) {
200            return self;
201        }
202        let dm = self.degree_multiple().max(rhs.degree_multiple());
203        Self::Sub {
204            x: Arc::new(self),
205            y: Arc::new(rhs),
206            degree_multiple: dm,
207        }
208    }
209
210    /// Negation with constant folding.
211    fn sym_neg(self) -> Self {
212        if let Some(&c) = self.as_const() {
213            return Self::Leaf(A::from_const(-c));
214        }
215        let dm = self.degree_multiple();
216        Self::Neg {
217            x: Arc::new(self),
218            degree_multiple: dm,
219        }
220    }
221
222    /// Multiplication with constant folding, zero-annihilation, and one-identity.
223    fn sym_mul(self, rhs: Self) -> Self {
224        if let (Some(&a), Some(&b)) = (self.as_const(), rhs.as_const()) {
225            return Self::Leaf(A::from_const(a * b));
226        }
227        if self.as_const().is_some_and(|c| c.is_zero())
228            || rhs.as_const().is_some_and(|c| c.is_zero())
229        {
230            return Self::Leaf(A::from_const(A::F::ZERO));
231        }
232        if self.as_const().is_some_and(|c| c.is_one()) {
233            return rhs;
234        }
235        if rhs.as_const().is_some_and(|c| c.is_one()) {
236            return self;
237        }
238        let dm = self.degree_multiple() + rhs.degree_multiple();
239        Self::Mul {
240            x: Arc::new(self),
241            y: Arc::new(rhs),
242            degree_multiple: dm,
243        }
244    }
245}
246
247impl<A: SymLeaf> PrimeCharacteristicRing for SymbolicExpr<A> {
248    type PrimeSubfield = <A::F as PrimeCharacteristicRing>::PrimeSubfield;
249
250    const ZERO: Self = Self::Leaf(A::ZERO);
251    const ONE: Self = Self::Leaf(A::ONE);
252    const TWO: Self = Self::Leaf(A::TWO);
253    const NEG_ONE: Self = Self::Leaf(A::NEG_ONE);
254
255    #[inline]
256    fn from_prime_subfield(f: Self::PrimeSubfield) -> Self {
257        Self::Leaf(A::from_const(A::F::from_prime_subfield(f)))
258    }
259}
260
261impl<A: SymLeaf> Dup for SymbolicExpr<A> {
262    #[inline(always)]
263    fn dup(&self) -> Self {
264        self.clone()
265    }
266}
267
268impl<A: SymLeaf> Default for SymbolicExpr<A> {
269    fn default() -> Self {
270        Self::ZERO
271    }
272}
273
274impl<A: SymLeaf, T: Into<Self>> ops::Add<T> for SymbolicExpr<A> {
275    type Output = Self;
276    fn add(self, rhs: T) -> Self {
277        self.sym_add(rhs.into())
278    }
279}
280
281impl<A: SymLeaf, T: Into<Self>> ops::Sub<T> for SymbolicExpr<A> {
282    type Output = Self;
283    fn sub(self, rhs: T) -> Self {
284        self.sym_sub(rhs.into())
285    }
286}
287
288impl<A: SymLeaf> ops::Neg for SymbolicExpr<A> {
289    type Output = Self;
290    fn neg(self) -> Self {
291        self.sym_neg()
292    }
293}
294
295impl<A: SymLeaf, T: Into<Self>> ops::Mul<T> for SymbolicExpr<A> {
296    type Output = Self;
297    fn mul(self, rhs: T) -> Self {
298        self.sym_mul(rhs.into())
299    }
300}
301
302impl<A: SymLeaf, T: Into<Self>> ops::AddAssign<T> for SymbolicExpr<A> {
303    fn add_assign(&mut self, rhs: T) {
304        *self = self.clone() + rhs.into();
305    }
306}
307
308impl<A: SymLeaf, T: Into<Self>> ops::SubAssign<T> for SymbolicExpr<A> {
309    fn sub_assign(&mut self, rhs: T) {
310        *self = self.clone() - rhs.into();
311    }
312}
313
314impl<A: SymLeaf, T: Into<Self>> ops::MulAssign<T> for SymbolicExpr<A> {
315    fn mul_assign(&mut self, rhs: T) {
316        *self = self.clone() * rhs.into();
317    }
318}
319
320impl<A: SymLeaf, T: Into<Self>> Sum<T> for SymbolicExpr<A> {
321    fn sum<I: Iterator<Item = T>>(iter: I) -> Self {
322        iter.map(Into::into)
323            .reduce(|a, b| a + b)
324            .unwrap_or(Self::ZERO)
325    }
326}
327
328impl<A: SymLeaf, T: Into<Self>> Product<T> for SymbolicExpr<A> {
329    fn product<I: Iterator<Item = T>>(iter: I) -> Self {
330        iter.map(Into::into)
331            .reduce(|a, b| a * b)
332            .unwrap_or(Self::ONE)
333    }
334}
335
336impl<F: Field, T: Into<SymbolicExpression<F>>> ops::Add<T> for SymbolicVariable<F> {
337    type Output = SymbolicExpression<F>;
338    fn add(self, rhs: T) -> Self::Output {
339        Self::Output::from(self) + rhs.into()
340    }
341}
342
343impl<F: Field, T: Into<SymbolicExpression<F>>> ops::Sub<T> for SymbolicVariable<F> {
344    type Output = SymbolicExpression<F>;
345    fn sub(self, rhs: T) -> Self::Output {
346        Self::Output::from(self) - rhs.into()
347    }
348}
349
350impl<F: Field, T: Into<SymbolicExpression<F>>> ops::Mul<T> for SymbolicVariable<F> {
351    type Output = SymbolicExpression<F>;
352    fn mul(self, rhs: T) -> Self::Output {
353        Self::Output::from(self) * rhs.into()
354    }
355}
356
357impl<F: Field, EF: ExtensionField<F>, T: Into<SymbolicExpressionExt<F, EF>>> ops::Add<T>
358    for SymbolicVariableExt<F, EF>
359{
360    type Output = SymbolicExpressionExt<F, EF>;
361    fn add(self, rhs: T) -> Self::Output {
362        Self::Output::from(self) + rhs.into()
363    }
364}
365
366impl<F: Field, EF: ExtensionField<F>, T: Into<SymbolicExpressionExt<F, EF>>> ops::Sub<T>
367    for SymbolicVariableExt<F, EF>
368{
369    type Output = SymbolicExpressionExt<F, EF>;
370    fn sub(self, rhs: T) -> Self::Output {
371        Self::Output::from(self) - rhs.into()
372    }
373}
374
375impl<F: Field, EF: ExtensionField<F>, T: Into<SymbolicExpressionExt<F, EF>>> ops::Mul<T>
376    for SymbolicVariableExt<F, EF>
377{
378    type Output = SymbolicExpressionExt<F, EF>;
379    fn mul(self, rhs: T) -> Self::Output {
380        Self::Output::from(self) * rhs.into()
381    }
382}
383
384#[cfg(test)]
385mod tests {
386    use p3_baby_bear::BabyBear;
387    use p3_field::extension::BinomialExtensionField;
388
389    use super::*;
390    use crate::symbolic::expression::BaseLeaf;
391    use crate::symbolic::expression_ext::ExtLeaf;
392    use crate::symbolic::variable::{BaseEntry, ExtEntry};
393
394    type F = BabyBear;
395    type EF = BinomialExtensionField<BabyBear, 4>;
396
397    #[test]
398    fn symbolic_variable_add_produces_add_node() {
399        // Adding a variable and a non-zero constant creates an addition node.
400        let var = SymbolicVariable::<F>::new(BaseEntry::Main { offset: 0 }, 0);
401        let expr = SymbolicExpression::from(F::new(5));
402        let result = var + expr;
403        match result {
404            SymbolicExpr::Add {
405                x,
406                y,
407                degree_multiple,
408            } => {
409                assert_eq!(degree_multiple, 1);
410                assert!(matches!(
411                    x.as_ref(),
412                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
413                        if v.index == 0 && v.entry == BaseEntry::Main { offset: 0 }
414                ));
415                assert!(matches!(
416                    y.as_ref(),
417                    SymbolicExpr::Leaf(BaseLeaf::Constant(c)) if *c == F::new(5)
418                ));
419            }
420            _ => panic!("Expected an Add node"),
421        }
422    }
423
424    #[test]
425    fn symbolic_variable_sub_produces_sub_node() {
426        // Subtracting two variables creates a subtraction node.
427        let var = SymbolicVariable::<F>::new(BaseEntry::Main { offset: 0 }, 0);
428        let other = SymbolicExpression::Leaf(BaseLeaf::Variable(SymbolicVariable::new(
429            BaseEntry::Main { offset: 0 },
430            1,
431        )));
432        let result = var - other;
433        match result {
434            SymbolicExpr::Sub {
435                x,
436                y,
437                degree_multiple,
438            } => {
439                assert_eq!(degree_multiple, 1);
440                assert!(matches!(
441                    x.as_ref(),
442                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
443                        if v.index == 0 && v.entry == BaseEntry::Main { offset: 0 }
444                ));
445                assert!(matches!(
446                    y.as_ref(),
447                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
448                        if v.index == 1 && v.entry == BaseEntry::Main { offset: 0 }
449                ));
450            }
451            _ => panic!("Expected a Sub node"),
452        }
453    }
454
455    #[test]
456    fn symbolic_variable_mul_produces_mul_node() {
457        // Multiplying two variables creates a multiplication node with summed degree.
458        let var = SymbolicVariable::<F>::new(BaseEntry::Main { offset: 0 }, 0);
459        let other = SymbolicExpression::Leaf(BaseLeaf::Variable(SymbolicVariable::new(
460            BaseEntry::Main { offset: 0 },
461            1,
462        )));
463        let result = var * other;
464        match result {
465            SymbolicExpr::Mul {
466                x,
467                y,
468                degree_multiple,
469            } => {
470                assert_eq!(degree_multiple, 2);
471                assert!(matches!(
472                    x.as_ref(),
473                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
474                        if v.index == 0 && v.entry == BaseEntry::Main { offset: 0 }
475                ));
476                assert!(matches!(
477                    y.as_ref(),
478                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
479                        if v.index == 1 && v.entry == BaseEntry::Main { offset: 0 }
480                ));
481            }
482            _ => panic!("Expected a Mul node"),
483        }
484    }
485
486    #[test]
487    fn symbolic_variable_ext_add_produces_add_node() {
488        // Adding an extension variable and a non-zero constant creates an addition node.
489        let var = SymbolicVariableExt::<F, EF>::new(ExtEntry::Permutation { offset: 0 }, 0);
490        let expr = SymbolicExpressionExt::<F, EF>::from(F::new(3));
491        let result = var + expr;
492        match result {
493            SymbolicExpr::Add {
494                x,
495                y,
496                degree_multiple,
497            } => {
498                assert_eq!(degree_multiple, 1);
499                assert!(matches!(
500                    x.as_ref(),
501                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
502                        if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
503                ));
504                assert!(matches!(
505                    y.as_ref(),
506                    SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Constant(c))))
507                        if *c == F::new(3)
508                ));
509            }
510            _ => panic!("Expected an Add node"),
511        }
512    }
513
514    #[test]
515    fn symbolic_variable_ext_sub_produces_sub_node() {
516        // Subtracting two extension variables creates a subtraction node.
517        let var = SymbolicVariableExt::<F, EF>::new(ExtEntry::Permutation { offset: 0 }, 0);
518        let other = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
519            ExtEntry::Permutation { offset: 0 },
520            1,
521        ));
522        let result = var - other;
523        match result {
524            SymbolicExpr::Sub {
525                x,
526                y,
527                degree_multiple,
528            } => {
529                assert_eq!(degree_multiple, 1);
530                assert!(matches!(
531                    x.as_ref(),
532                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
533                        if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
534                ));
535                assert!(matches!(
536                    y.as_ref(),
537                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
538                        if v.index == 1 && v.entry == ExtEntry::Permutation { offset: 0 }
539                ));
540            }
541            _ => panic!("Expected a Sub node"),
542        }
543    }
544
545    #[test]
546    fn symbolic_variable_ext_mul_produces_mul_node() {
547        // Multiplying two extension variables creates a multiplication node with summed degree.
548        let var = SymbolicVariableExt::<F, EF>::new(ExtEntry::Permutation { offset: 0 }, 0);
549        let other = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
550            ExtEntry::Permutation { offset: 0 },
551            1,
552        ));
553        let result = var * other;
554        match result {
555            SymbolicExpr::Mul {
556                x,
557                y,
558                degree_multiple,
559            } => {
560                assert_eq!(degree_multiple, 2);
561                assert!(matches!(
562                    x.as_ref(),
563                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
564                        if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
565                ));
566                assert!(matches!(
567                    y.as_ref(),
568                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
569                        if v.index == 1 && v.entry == ExtEntry::Permutation { offset: 0 }
570                ));
571            }
572            _ => panic!("Expected a Mul node"),
573        }
574    }
575}
p3_air/symbolic/mod.rs

p3_air/symbolic/
mod.rs
