p3_air/symbolic/
expression_ext.rs

1use alloc::sync::Arc;
2
3use p3_field::extension::{
4    BinomialExtensionField, Complex, CubicTrinomialExtensionField, QuinticTrinomialExtensionField,
5};
6use p3_field::{Algebra, ExtensionField, Field, PrimeCharacteristicRing};
7use serde::{Deserialize, Serialize};
8
9use crate::symbolic::expression::BaseLeaf;
10use crate::symbolic::variable::SymbolicVariableExt;
11use crate::symbolic::{SymLeaf, SymbolicExpr, SymbolicExpression, SymbolicVariable};
12
13/// Leaf nodes for extension-field symbolic expressions.
14///
15/// These represent the atomic building blocks of extension-field AIR constraints:
16/// lifted base-field sub-trees, extension-field variables, and extension-field constants.
17#[derive(Clone, Debug, Serialize, Deserialize)]
18pub enum ExtLeaf<F, EF> {
19    /// A lifted base-field expression (entire base sub-tree preserved).
20    Base(SymbolicExpression<F>),
21
22    /// An extension-field variable (permutation column or challenge).
23    ExtVariable(SymbolicVariableExt<F, EF>),
24
25    /// An extension-field constant.
26    ExtConstant(EF),
27}
28
29/// A symbolic expression tree for extension-field AIR constraints.
30///
31/// This is a type alias for the generic [`SymbolicExpr`] parameterized with
32/// extension-field [`ExtLeaf`] nodes.
33pub type SymbolicExpressionExt<F, EF> = SymbolicExpr<ExtLeaf<F, EF>>;
34
35impl<F: Field, EF: ExtensionField<F>> SymLeaf for ExtLeaf<F, EF> {
36    type F = F;
37
38    const ZERO: Self = Self::Base(SymbolicExpression::ZERO);
39    const ONE: Self = Self::Base(SymbolicExpression::ONE);
40    const TWO: Self = Self::Base(SymbolicExpression::TWO);
41    const NEG_ONE: Self = Self::Base(SymbolicExpression::NEG_ONE);
42
43    fn degree_multiple(&self) -> usize {
44        match self {
45            Self::Base(e) => e.degree_multiple(),
46            Self::ExtVariable(v) => v.degree_multiple(),
47            Self::ExtConstant(_) => 0,
48        }
49    }
50
51    fn poly_degree(&self, trace_len: usize, periodic_periods: &[usize]) -> usize {
52        match self {
53            Self::Base(e) => e.poly_degree(trace_len, periodic_periods),
54            Self::ExtVariable(v) => v.poly_degree(trace_len),
55            Self::ExtConstant(_) => 0,
56        }
57    }
58
59    fn as_const(&self) -> Option<&F> {
60        match self {
61            Self::Base(SymbolicExpression::Leaf(BaseLeaf::Constant(c))) => Some(c),
62            Self::ExtConstant(ef) if ef.is_in_basefield() => {
63                Some(&ef.as_basis_coefficients_slice()[0])
64            }
65            _ => None,
66        }
67    }
68
69    fn from_const(c: F) -> Self {
70        Self::Base(SymbolicExpression::from(c))
71    }
72}
73
74impl<F: Field, EF> SymbolicExpressionExt<F, EF> {
75    /// Try to lower this extension expression to a base-field expression.
76    ///
77    /// Returns `None` if the tree contains any extension-only nodes
78    /// ([`ExtVariable`](ExtLeaf::ExtVariable) or [`ExtConstant`](ExtLeaf::ExtConstant)).
79    pub fn to_base(&self) -> Option<SymbolicExpression<F>> {
80        match self {
81            Self::Leaf(ExtLeaf::Base(e)) => Some(e.clone()),
82            Self::Leaf(ExtLeaf::ExtVariable(_) | ExtLeaf::ExtConstant(_)) => None,
83            Self::Add {
84                x,
85                y,
86                degree_multiple,
87            } => Some(SymbolicExpr::Add {
88                x: Arc::new(x.to_base()?),
89                y: Arc::new(y.to_base()?),
90                degree_multiple: *degree_multiple,
91            }),
92            Self::Sub {
93                x,
94                y,
95                degree_multiple,
96            } => Some(SymbolicExpr::Sub {
97                x: Arc::new(x.to_base()?),
98                y: Arc::new(y.to_base()?),
99                degree_multiple: *degree_multiple,
100            }),
101            Self::Neg { x, degree_multiple } => Some(SymbolicExpr::Neg {
102                x: Arc::new(x.to_base()?),
103                degree_multiple: *degree_multiple,
104            }),
105            Self::Mul {
106                x,
107                y,
108                degree_multiple,
109            } => Some(SymbolicExpr::Mul {
110                x: Arc::new(x.to_base()?),
111                y: Arc::new(y.to_base()?),
112                degree_multiple: *degree_multiple,
113            }),
114        }
115    }
116}
117
118impl<F: Field, EF> From<SymbolicExpression<F>> for SymbolicExpressionExt<F, EF> {
119    fn from(expr: SymbolicExpression<F>) -> Self {
120        Self::Leaf(ExtLeaf::Base(expr))
121    }
122}
123
124impl<F: Field, EF> From<SymbolicVariable<F>> for SymbolicExpressionExt<F, EF> {
125    fn from(var: SymbolicVariable<F>) -> Self {
126        Self::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Variable(var))))
127    }
128}
129
130impl<F, EF> From<SymbolicVariableExt<F, EF>> for SymbolicExpressionExt<F, EF> {
131    fn from(var: SymbolicVariableExt<F, EF>) -> Self {
132        Self::Leaf(ExtLeaf::ExtVariable(var))
133    }
134}
135
136impl<F: Field, EF> From<F> for SymbolicExpressionExt<F, EF> {
137    fn from(f: F) -> Self {
138        Self::Leaf(ExtLeaf::Base(SymbolicExpression::from(f)))
139    }
140}
141
142/// Concrete [`From`] for [`BinomialExtensionField`] constants.
143///
144/// This avoids overlap with [`From<F>`] when `EF = F`, since
145/// [`BinomialExtensionField<F, D>`] is always a distinct type from `F`.
146impl<F, const D: usize> From<BinomialExtensionField<F, D>>
147    for SymbolicExpressionExt<F, BinomialExtensionField<F, D>>
148where
149    F: Field,
150    BinomialExtensionField<F, D>: ExtensionField<F>,
151{
152    fn from(ef: BinomialExtensionField<F, D>) -> Self {
153        Self::Leaf(ExtLeaf::ExtConstant(ef))
154    }
155}
156
157impl<F: Field, EF: ExtensionField<F>> Algebra<F> for SymbolicExpressionExt<F, EF> {}
158
159impl<F: Field, EF: ExtensionField<F>> Algebra<SymbolicExpression<F>>
160    for SymbolicExpressionExt<F, EF>
161{
162}
163
164impl<F: Field, EF: ExtensionField<F>> Algebra<SymbolicVariable<F>>
165    for SymbolicExpressionExt<F, EF>
166{
167}
168
169impl<F: Field, EF: ExtensionField<F>> Algebra<SymbolicVariableExt<F, EF>>
170    for SymbolicExpressionExt<F, EF>
171{
172}
173
174/// Concrete [`Algebra`] for [`BinomialExtensionField`] — avoids overlap with `Algebra<F>` when `EF = F`.
175impl<F: Field, const D: usize> Algebra<BinomialExtensionField<F, D>>
176    for SymbolicExpressionExt<F, BinomialExtensionField<F, D>>
177where
178    BinomialExtensionField<F, D>: ExtensionField<F>,
179{
180}
181
182impl<F: Field> From<CubicTrinomialExtensionField<F>>
183    for SymbolicExpressionExt<F, CubicTrinomialExtensionField<F>>
184where
185    CubicTrinomialExtensionField<F>: ExtensionField<F>,
186{
187    fn from(ef: CubicTrinomialExtensionField<F>) -> Self {
188        Self::Leaf(ExtLeaf::ExtConstant(ef))
189    }
190}
191
192impl<F: Field> From<QuinticTrinomialExtensionField<F>>
193    for SymbolicExpressionExt<F, QuinticTrinomialExtensionField<F>>
194where
195    QuinticTrinomialExtensionField<F>: ExtensionField<F>,
196{
197    fn from(ef: QuinticTrinomialExtensionField<F>) -> Self {
198        Self::Leaf(ExtLeaf::ExtConstant(ef))
199    }
200}
201
202/// Concrete [`From`] for a degree-4 complex tower `BinomialExtensionField<Complex<F>, 2>`.
203///
204/// The symbolic base is `F` while the binomial's base parameter is `Complex<F>`, so
205/// the generic [`BinomialExtensionField<F, D>`] impl above does not cover it.
206impl<F: Field> From<BinomialExtensionField<Complex<F>, 2>>
207    for SymbolicExpressionExt<F, BinomialExtensionField<Complex<F>, 2>>
208where
209    BinomialExtensionField<Complex<F>, 2>: ExtensionField<F>,
210{
211    fn from(ef: BinomialExtensionField<Complex<F>, 2>) -> Self {
212        Self::Leaf(ExtLeaf::ExtConstant(ef))
213    }
214}
215
216/// Concrete [`Algebra`] for [`CubicTrinomialExtensionField`] — avoids overlap with `Algebra<F>`.
217impl<F: Field> Algebra<CubicTrinomialExtensionField<F>>
218    for SymbolicExpressionExt<F, CubicTrinomialExtensionField<F>>
219where
220    CubicTrinomialExtensionField<F>: ExtensionField<F>,
221{
222}
223
224/// Concrete [`Algebra`] for [`QuinticTrinomialExtensionField`] — avoids overlap with `Algebra<F>`.
225impl<F: Field> Algebra<QuinticTrinomialExtensionField<F>>
226    for SymbolicExpressionExt<F, QuinticTrinomialExtensionField<F>>
227where
228    QuinticTrinomialExtensionField<F>: ExtensionField<F>,
229{
230}
231
232/// Concrete [`Algebra`] for a degree-4 complex tower `BinomialExtensionField<Complex<F>, 2>` —
233/// avoids overlap with `Algebra<F>` and with the generic binomial impl, whose base parameter
234/// matches the symbolic base.
235impl<F: Field> Algebra<BinomialExtensionField<Complex<F>, 2>>
236    for SymbolicExpressionExt<F, BinomialExtensionField<Complex<F>, 2>>
237where
238    BinomialExtensionField<Complex<F>, 2>: ExtensionField<F>,
239{
240}
241
242#[cfg(test)]
243mod tests {
244    use p3_baby_bear::BabyBear;
245    use p3_field::extension::BinomialExtensionField;
246    use p3_field::{BasedVectorSpace, PrimeCharacteristicRing};
247    use p3_mersenne_31::{Mersenne31, QM31};
248
249    use super::*;
250    use crate::symbolic::SymbolicExpr;
251    use crate::symbolic::variable::{BaseEntry, ExtEntry};
252
253    type F = BabyBear;
254    type EF = BinomialExtensionField<BabyBear, 4>;
255
256    #[test]
257    fn ext_leaf_degree_multiple_base_variable() {
258        // A base leaf with a trace variable inside has degree 1.
259        let var = SymbolicVariable::<F>::new(BaseEntry::Main { offset: 0 }, 0);
260        let leaf = ExtLeaf::<F, EF>::Base(SymbolicExpression::from(var));
261        assert_eq!(leaf.degree_multiple(), 1);
262    }
263
264    #[test]
265    fn ext_leaf_degree_multiple_base_constant() {
266        // A base leaf with a constant inside has degree 0.
267        let leaf = ExtLeaf::<F, EF>::Base(SymbolicExpression::from(F::new(42)));
268        assert_eq!(leaf.degree_multiple(), 0);
269    }
270
271    #[test]
272    fn ext_leaf_degree_multiple_ext_variable() {
273        // A permutation variable has degree 1.
274        let var = SymbolicVariableExt::<F, EF>::new(ExtEntry::Permutation { offset: 0 }, 0);
275        let leaf = ExtLeaf::ExtVariable(var);
276        assert_eq!(leaf.degree_multiple(), 1);
277    }
278
279    #[test]
280    fn ext_leaf_degree_multiple_ext_variable_challenge() {
281        // A challenge variable has degree 0.
282        let var = SymbolicVariableExt::<F, EF>::new(ExtEntry::Challenge, 0);
283        let leaf = ExtLeaf::ExtVariable(var);
284        assert_eq!(leaf.degree_multiple(), 0);
285    }
286
287    #[test]
288    fn ext_leaf_degree_multiple_ext_constant() {
289        // An extension constant always has degree 0.
290        let leaf = ExtLeaf::<F, EF>::ExtConstant(EF::ONE);
291        assert_eq!(leaf.degree_multiple(), 0);
292    }
293
294    #[test]
295    fn ext_leaf_as_const_base_constant() {
296        // A base constant leaf can be viewed as a field constant.
297        let leaf = ExtLeaf::<F, EF>::Base(SymbolicExpression::from(F::new(7)));
298        assert_eq!(leaf.as_const(), Some(&F::new(7)));
299    }
300
301    #[test]
302    fn ext_leaf_as_const_base_variable() {
303        // A base variable leaf is not a constant.
304        let var = SymbolicVariable::<F>::new(BaseEntry::Main { offset: 0 }, 0);
305        let leaf = ExtLeaf::<F, EF>::Base(SymbolicExpression::from(var));
306        assert!(leaf.as_const().is_none());
307    }
308
309    #[test]
310    fn ext_leaf_as_const_ext_variable() {
311        // An extension variable leaf is not a constant.
312        let var = SymbolicVariableExt::<F, EF>::new(ExtEntry::Permutation { offset: 0 }, 0);
313        let leaf = ExtLeaf::ExtVariable(var);
314        assert!(leaf.as_const().is_none());
315    }
316
317    #[test]
318    fn ext_leaf_as_const_ext_constant_in_basefield() {
319        // An extension constant that lies in the base field is recognized as a constant.
320        let leaf = ExtLeaf::<F, EF>::ExtConstant(EF::ONE);
321        assert_eq!(leaf.as_const(), Some(&F::ONE));
322    }
323
324    #[test]
325    fn ext_leaf_as_const_ext_constant_zero() {
326        // The extension zero element is recognized as the base zero.
327        let leaf = ExtLeaf::<F, EF>::ExtConstant(EF::ZERO);
328        assert_eq!(leaf.as_const(), Some(&F::ZERO));
329    }
330
331    #[test]
332    fn ext_leaf_as_const_ext_constant_not_in_basefield() {
333        // An extension constant with non-zero higher coefficients is not a base constant.
334        let ef_val = EF::from_basis_coefficients_fn(|i| if i == 1 { F::ONE } else { F::ZERO });
335        let leaf = ExtLeaf::<F, EF>::ExtConstant(ef_val);
336        assert!(leaf.as_const().is_none());
337    }
338
339    #[test]
340    fn ext_leaf_from_const() {
341        // Creating a leaf from a base-field value produces a constant.
342        let leaf = ExtLeaf::<F, EF>::from_const(F::new(13));
343        assert_eq!(leaf.as_const(), Some(&F::new(13)));
344    }
345
346    #[test]
347    fn to_base_leaf_base() {
348        // A base-only leaf can be lowered to a base expression.
349        let base_expr = SymbolicExpression::from(F::new(5));
350        let ext_expr = SymbolicExpressionExt::<F, EF>::from(base_expr);
351        let lowered = ext_expr.to_base();
352
353        assert!(lowered.is_some());
354        assert!(matches!(
355            lowered.unwrap(),
356            SymbolicExpr::Leaf(BaseLeaf::Constant(c)) if c == F::new(5)
357        ));
358    }
359
360    #[test]
361    fn to_base_leaf_ext_variable() {
362        // An extension variable cannot be lowered to base.
363        let var = SymbolicVariableExt::<F, EF>::new(ExtEntry::Permutation { offset: 0 }, 0);
364        let ext_expr = SymbolicExpressionExt::<F, EF>::from(var);
365        assert!(ext_expr.to_base().is_none());
366    }
367
368    #[test]
369    fn to_base_leaf_ext_constant() {
370        // An extension constant cannot be lowered to base.
371        let ext_expr = SymbolicExpressionExt::<F, EF>::Leaf(ExtLeaf::ExtConstant(EF::TWO));
372        assert!(ext_expr.to_base().is_none());
373    }
374
375    #[test]
376    fn to_base_add_of_base_exprs() {
377        // A sum of two base-only expressions can be lowered.
378        let a = SymbolicExpressionExt::<F, EF>::from(F::new(3));
379        let b = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
380            BaseEntry::Main { offset: 0 },
381            0,
382        ));
383        let sum = a + b;
384        let lowered = sum.to_base();
385
386        match lowered {
387            Some(SymbolicExpr::Add {
388                x,
389                y,
390                degree_multiple,
391            }) => {
392                assert_eq!(degree_multiple, 1);
393                assert!(matches!(
394                    x.as_ref(),
395                    SymbolicExpr::Leaf(BaseLeaf::Constant(c)) if *c == F::new(3)
396                ));
397                assert!(matches!(
398                    y.as_ref(),
399                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
400                        if v.index == 0 && v.entry == BaseEntry::Main { offset: 0 }
401                ));
402            }
403            _ => panic!("Expected a lowered Add node"),
404        }
405    }
406
407    #[test]
408    fn to_base_add_with_ext_child_returns_none() {
409        // A sum with one extension-only child cannot be lowered.
410        let base = SymbolicExpressionExt::<F, EF>::from(F::new(3));
411        let ext_var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
412            ExtEntry::Permutation { offset: 0 },
413            0,
414        ));
415        let sum = base + ext_var;
416        assert!(sum.to_base().is_none());
417    }
418
419    #[test]
420    fn to_base_sub_of_base_exprs() {
421        // A difference of two base-only expressions can be lowered.
422        let a = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
423            BaseEntry::Main { offset: 0 },
424            0,
425        ));
426        let b = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
427            BaseEntry::Main { offset: 0 },
428            1,
429        ));
430        let diff = a - b;
431        let lowered = diff.to_base();
432
433        match lowered {
434            Some(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 lowered Sub node"),
452        }
453    }
454
455    #[test]
456    fn to_base_sub_with_ext_child_returns_none() {
457        // A difference with an extension-only child cannot be lowered.
458        let base = SymbolicExpressionExt::<F, EF>::from(F::new(5));
459        let ext_var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
460            ExtEntry::Challenge,
461            0,
462        ));
463        let diff = base - ext_var;
464        assert!(diff.to_base().is_none());
465    }
466
467    #[test]
468    fn to_base_neg_of_base_expr() {
469        // Negation of a base-only expression can be lowered.
470        let var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
471            BaseEntry::Main { offset: 0 },
472            0,
473        ));
474        let neg = -var;
475        let lowered = neg.to_base();
476
477        match lowered {
478            Some(SymbolicExpr::Neg { x, degree_multiple }) => {
479                assert_eq!(degree_multiple, 1);
480                assert!(matches!(
481                    x.as_ref(),
482                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
483                        if v.index == 0 && v.entry == BaseEntry::Main { offset: 0 }
484                ));
485            }
486            _ => panic!("Expected a lowered Neg node"),
487        }
488    }
489
490    #[test]
491    fn to_base_mul_of_base_exprs() {
492        // A product of two base-only expressions can be lowered.
493        let a = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
494            BaseEntry::Main { offset: 0 },
495            0,
496        ));
497        let b = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
498            BaseEntry::Main { offset: 0 },
499            1,
500        ));
501        let prod = a * b;
502        let lowered = prod.to_base();
503
504        match lowered {
505            Some(SymbolicExpr::Mul {
506                x,
507                y,
508                degree_multiple,
509            }) => {
510                assert_eq!(degree_multiple, 2);
511                assert!(matches!(
512                    x.as_ref(),
513                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
514                        if v.index == 0 && v.entry == BaseEntry::Main { offset: 0 }
515                ));
516                assert!(matches!(
517                    y.as_ref(),
518                    SymbolicExpr::Leaf(BaseLeaf::Variable(v))
519                        if v.index == 1 && v.entry == BaseEntry::Main { offset: 0 }
520                ));
521            }
522            _ => panic!("Expected a lowered Mul node"),
523        }
524    }
525
526    #[test]
527    fn to_base_mul_with_ext_child_returns_none() {
528        // A product with an extension-only child cannot be lowered.
529        let base = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
530            BaseEntry::Main { offset: 0 },
531            0,
532        ));
533        let ext_var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
534            ExtEntry::Permutation { offset: 0 },
535            0,
536        ));
537        let prod = base * ext_var;
538        assert!(prod.to_base().is_none());
539    }
540
541    #[test]
542    fn from_symbolic_expression() {
543        // Converting a base expression lifts it into a base leaf.
544        let base_expr = SymbolicExpression::from(F::new(99));
545        let ext_expr = SymbolicExpressionExt::<F, EF>::from(base_expr);
546        assert!(matches!(
547            ext_expr,
548            SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Constant(c)))) if c == F::new(99)
549        ));
550    }
551
552    #[test]
553    fn from_symbolic_variable() {
554        // Converting a base variable lifts it into a base leaf.
555        let var = SymbolicVariable::<F>::new(BaseEntry::Main { offset: 0 }, 2);
556        let ext_expr = SymbolicExpressionExt::<F, EF>::from(var);
557        assert!(matches!(
558            ext_expr,
559            SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Variable(v))))
560                if v.index == 2 && v.entry == BaseEntry::Main { offset: 0 }
561        ));
562    }
563
564    #[test]
565    fn from_symbolic_variable_ext() {
566        // Converting an extension variable produces an extension variable leaf.
567        let var = SymbolicVariableExt::<F, EF>::new(ExtEntry::Permutation { offset: 1 }, 3);
568        let ext_expr = SymbolicExpressionExt::<F, EF>::from(var);
569        assert!(matches!(
570            ext_expr,
571            SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
572                if v.index == 3 && v.entry == ExtEntry::Permutation { offset: 1 }
573        ));
574    }
575
576    #[test]
577    fn from_base_field() {
578        // Converting a base field element produces a base constant leaf.
579        let ext_expr = SymbolicExpressionExt::<F, EF>::from(F::new(42));
580        assert!(matches!(
581            ext_expr,
582            SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Constant(c)))) if c == F::new(42)
583        ));
584    }
585
586    #[test]
587    fn from_binomial_extension_field() {
588        // Converting an extension field element produces an extension constant leaf.
589        let ef_val = EF::ONE + EF::ONE;
590        let ext_expr = SymbolicExpressionExt::<F, EF>::from(ef_val);
591        assert!(matches!(
592            ext_expr,
593            SymbolicExpr::Leaf(ExtLeaf::ExtConstant(c)) if c == ef_val
594        ));
595    }
596
597    #[test]
598    fn ext_add_constant_folding() {
599        // Two base constants fold into one on addition.
600        let a = SymbolicExpressionExt::<F, EF>::from(F::new(3));
601        let b = SymbolicExpressionExt::<F, EF>::from(F::new(4));
602        let result = a + b;
603        assert!(matches!(
604            result,
605            SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Constant(c)))) if c == F::new(7)
606        ));
607    }
608
609    #[test]
610    fn ext_sub_constant_folding() {
611        // Two base constants fold into one on subtraction.
612        let a = SymbolicExpressionExt::<F, EF>::from(F::new(10));
613        let b = SymbolicExpressionExt::<F, EF>::from(F::new(4));
614        let result = a - b;
615        assert!(matches!(
616            result,
617            SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Constant(c)))) if c == F::new(6)
618        ));
619    }
620
621    #[test]
622    fn ext_mul_constant_folding() {
623        // Two base constants fold into one on multiplication.
624        let a = SymbolicExpressionExt::<F, EF>::from(F::new(3));
625        let b = SymbolicExpressionExt::<F, EF>::from(F::new(5));
626        let result = a * b;
627        assert!(matches!(
628            result,
629            SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Constant(c)))) if c == F::new(15)
630        ));
631    }
632
633    #[test]
634    fn ext_add_variables_degree_tracking() {
635        // Adding two degree-1 variables gives degree 1 (the max).
636        let a = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
637            ExtEntry::Permutation { offset: 0 },
638            0,
639        ));
640        let b = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
641            ExtEntry::Permutation { offset: 0 },
642            1,
643        ));
644        let result = a + b;
645
646        match result {
647            SymbolicExpr::Add {
648                x,
649                y,
650                degree_multiple,
651            } => {
652                assert_eq!(degree_multiple, 1);
653                assert!(matches!(
654                    x.as_ref(),
655                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
656                        if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
657                ));
658                assert!(matches!(
659                    y.as_ref(),
660                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
661                        if v.index == 1 && v.entry == ExtEntry::Permutation { offset: 0 }
662                ));
663            }
664            _ => panic!("Expected an Add node"),
665        }
666    }
667
668    #[test]
669    fn ext_mul_variables_degree_tracking() {
670        // Multiplying two degree-1 variables gives degree 2 (the sum).
671        let a = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
672            ExtEntry::Permutation { offset: 0 },
673            0,
674        ));
675        let b = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
676            ExtEntry::Permutation { offset: 0 },
677            1,
678        ));
679        let result = a * b;
680
681        match result {
682            SymbolicExpr::Mul {
683                x,
684                y,
685                degree_multiple,
686            } => {
687                assert_eq!(degree_multiple, 2);
688                assert!(matches!(
689                    x.as_ref(),
690                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
691                        if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
692                ));
693                assert!(matches!(
694                    y.as_ref(),
695                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
696                        if v.index == 1 && v.entry == ExtEntry::Permutation { offset: 0 }
697                ));
698            }
699            _ => panic!("Expected a Mul node"),
700        }
701    }
702
703    #[test]
704    fn ext_constant_zero_mul_folds_to_zero() {
705        // Multiplying by the extension zero folds to the zero constant.
706        let var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
707            ExtEntry::Permutation { offset: 0 },
708            0,
709        ));
710        let zero = SymbolicExpressionExt::<F, EF>::from(EF::ZERO);
711        let result = var * zero;
712        assert!(matches!(
713            result,
714            SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Constant(c)))) if c == F::ZERO
715        ));
716    }
717
718    #[test]
719    fn ext_constant_one_mul_folds_to_identity() {
720        // Multiplying by the extension one folds to the other operand.
721        let var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
722            ExtEntry::Permutation { offset: 0 },
723            0,
724        ));
725        let one = SymbolicExpressionExt::<F, EF>::from(EF::ONE);
726        let result = var * one;
727        assert!(matches!(
728            result,
729            SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
730                if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
731        ));
732    }
733
734    #[test]
735    fn ext_constant_zero_add_folds_to_identity() {
736        // Adding the extension zero folds to the other operand.
737        let var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
738            ExtEntry::Permutation { offset: 0 },
739            0,
740        ));
741        let zero = SymbolicExpressionExt::<F, EF>::from(EF::ZERO);
742        let result = zero + var;
743        assert!(matches!(
744            result,
745            SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
746                if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
747        ));
748    }
749
750    #[test]
751    fn ext_constant_zero_sub_folds_to_neg() {
752        // Subtracting from the extension zero folds to negation.
753        let var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
754            ExtEntry::Permutation { offset: 0 },
755            0,
756        ));
757        let zero = SymbolicExpressionExt::<F, EF>::from(EF::ZERO);
758        let result = zero - var;
759        match result {
760            SymbolicExpr::Neg { x, degree_multiple } => {
761                assert_eq!(degree_multiple, 1);
762                assert!(matches!(
763                    x.as_ref(),
764                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
765                        if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
766                ));
767            }
768            _ => panic!("Expected a Neg node"),
769        }
770    }
771
772    #[test]
773    fn ext_constant_not_in_basefield_no_folding() {
774        // A non-base-field extension constant does not fold with multiplication.
775        let var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
776            ExtEntry::Permutation { offset: 0 },
777            0,
778        ));
779        let non_base = SymbolicExpressionExt::<F, EF>::from(EF::from_basis_coefficients_fn(|i| {
780            if i == 1 { F::ONE } else { F::ZERO }
781        }));
782        let result = var * non_base;
783        match result {
784            SymbolicExpr::Mul {
785                x,
786                y,
787                degree_multiple,
788            } => {
789                assert_eq!(degree_multiple, 1);
790                assert!(matches!(
791                    x.as_ref(),
792                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
793                        if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
794                ));
795                assert!(matches!(
796                    y.as_ref(),
797                    SymbolicExpr::Leaf(ExtLeaf::ExtConstant(_))
798                ));
799            }
800            _ => panic!("Expected a Mul node since the constant is not in the base field"),
801        }
802    }
803
804    #[test]
805    fn ext_mixed_base_and_ext_arithmetic() {
806        // Mixing a base variable with an extension variable in a sum.
807        let base_var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
808            BaseEntry::Main { offset: 0 },
809            0,
810        ));
811        let ext_var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
812            ExtEntry::Permutation { offset: 0 },
813            0,
814        ));
815        let result = base_var + ext_var;
816
817        match &result {
818            SymbolicExpr::Add {
819                x,
820                y,
821                degree_multiple,
822            } => {
823                assert_eq!(*degree_multiple, 1);
824                assert!(matches!(
825                    x.as_ref(),
826                    SymbolicExpr::Leaf(ExtLeaf::Base(SymbolicExpr::Leaf(BaseLeaf::Variable(v))))
827                        if v.index == 0 && v.entry == BaseEntry::Main { offset: 0 }
828                ));
829                assert!(matches!(
830                    y.as_ref(),
831                    SymbolicExpr::Leaf(ExtLeaf::ExtVariable(v))
832                        if v.index == 0 && v.entry == ExtEntry::Permutation { offset: 0 }
833                ));
834            }
835            _ => panic!("Expected an Add node"),
836        }
837
838        // The mixed result cannot be lowered to base.
839        assert!(result.to_base().is_none());
840    }
841
842    #[test]
843    fn serde_round_trip_preserves_extension_constraint() {
844        // A constraint over all extension leaf kinds and a lifted base sub-tree:
845        //   perm[0]·challenge - ext_const + base_var
846        let perm = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
847            ExtEntry::Permutation { offset: 0 },
848            0,
849        ));
850        let challenge = SymbolicExpressionExt::<F, EF>::from(SymbolicVariableExt::<F, EF>::new(
851            ExtEntry::Challenge,
852            0,
853        ));
854        let ext_const = SymbolicExpressionExt::<F, EF>::from(EF::from_basis_coefficients_fn(|i| {
855            if i == 1 { F::ONE } else { F::ZERO }
856        }));
857        let base_var = SymbolicExpressionExt::<F, EF>::from(SymbolicVariable::<F>::new(
858            BaseEntry::Main { offset: 0 },
859            0,
860        ));
861
862        let expr = perm * challenge - ext_const + base_var;
863
864        let json = serde_json::to_string(&expr).unwrap();
865        let decoded: SymbolicExpressionExt<F, EF> = serde_json::from_str(&json).unwrap();
866
867        // Structural equality: the decoded tree re-serializes identically.
868        assert_eq!(serde_json::to_string(&decoded).unwrap(), json);
869        assert_eq!(decoded.degree_multiple(), expr.degree_multiple());
870    }
871
872    #[test]
873    fn complex_tower_extension_constant_lowers_to_leaf() {
874        // `QM31 = BinomialExtensionField<Complex<Mersenne31>, 2>` is a degree-4 tower whose
875        // binomial base parameter (`Complex<Mersenne31>`) differs from the symbolic base
876        // (`Mersenne31`), so it needs the dedicated complex-tower impls.
877        fn assert_algebra<A: Algebra<B>, B>() {}
878        assert_algebra::<SymbolicExpressionExt<Mersenne31, QM31>, QM31>();
879
880        let expr = SymbolicExpressionExt::<Mersenne31, QM31>::from(QM31::ONE);
881        match expr {
882            SymbolicExpressionExt::Leaf(ExtLeaf::ExtConstant(c)) => assert_eq!(c, QM31::ONE),
883            _ => panic!("Expected an ExtConstant leaf"),
884        }
885    }
886}
p3_air/symbolic/expression_ext.rs

p3_air/symbolic/
expression_ext.rs
