Formal polynomial evaluator.

crates/prover/src/constraint_framework/mod.rs

@@ -7,6 +7,8 @@ mod info;
 pub mod logup;
 mod point;
 pub mod poly;
+pub mod poly_eval;
+pub mod rational;
 mod simd_domain;
 
 use std::array;

crates/prover/src/constraint_framework/poly.rs

@@ -1,6 +1,6 @@
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
-use std::ops::{Add, AddAssign, Mul, Neg};
+use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub};
 
 use itertools::Itertools;
 use num_traits::{One, Zero};
@@ -21,18 +21,37 @@ impl Monomial {
     }
 
     fn default() -> Monomial {
-        Monomial {
-            vars: [(0, 0)].into(),
-        }
+        Monomial { vars: [].into() }
+    }
+}
+
+impl One for Monomial {
+    fn one() -> Self {
+        Monomial { vars: [].into() }
     }
 }
 
 /// A polynomial consists of a list of monomials with coefficients.
 #[derive(Debug, Hash, PartialEq, PartialOrd, Eq, Ord, Clone)]
-struct Polynomial<F: From<BaseField>> {
+pub struct Polynomial<F: From<BaseField>> {
     monomials: BTreeMap<Monomial, F>,
 }
 
+impl<F: From<BaseField> + One> Polynomial<F> {
+    /// Returns the polynomial x_ind.
+    pub fn from_var_index(ind: usize) -> Polynomial<F> {
+        Self {
+            monomials: [(
+                Monomial {
+                    vars: [(ind, 1)].into(),
+                },
+                F::one(),
+            )]
+            .into(),
+        }
+    }
+}
+
 impl<F: One + From<BaseField>> From<Monomial> for Polynomial<F> {
     fn from(monomial: Monomial) -> Self {
         Self {
@@ -72,6 +91,16 @@ impl<F: Zero + Add + Clone + From<BaseField>> Add for Polynomial<F> {
     }
 }
 
+impl<F> Sub for Polynomial<F>
+where
+    F: Zero + Add + Clone + From<BaseField> + Neg<Output = F>,
+{
+    type Output = Self;
+    fn sub(self, rhs: Self) -> Self {
+        self + (-rhs)
+    }
+}
+
 #[allow(clippy::suspicious_arithmetic_impl)]
 impl Mul for Monomial {
     type Output = Self;
@@ -126,6 +155,16 @@ where
     }
 }
 
+impl<F, G> MulAssign<G> for Polynomial<F>
+where
+    F: Clone + Mul<Output = F> + Add<Output = F> + AddAssign + From<BaseField>,
+    G: Into<Polynomial<F>>,
+{
+    fn mul_assign(&mut self, rhs: G) {
+        *self = self.clone() * rhs.into();
+    }
+}
+
 impl<F> Mul<SecureField> for Polynomial<F>
 where
     F: Clone
@@ -157,6 +196,14 @@ impl<F: Zero + Clone + From<BaseField>> Zero for Polynomial<F> {
     }
 }
 
+impl<F: One + From<BaseField> + AddAssign + Clone + Add<Output = F>> One for Polynomial<F> {
+    fn one() -> Self {
+        Self {
+            monomials: [(Monomial::default(), F::one())].into(),
+        }
+    }
+}
+
 impl<F: Neg<Output = F> + From<BaseField>> Neg for Polynomial<F> {
     type Output = Self;
     fn neg(self) -> Self {
@@ -180,6 +227,32 @@ where
     }
 }
 
+impl From<Polynomial<BaseField>> for Polynomial<SecureField> {
+    fn from(poly: Polynomial<BaseField>) -> Self {
+        Self {
+            monomials: poly
+                .monomials
+                .into_iter()
+                .map(|(m, c)| (m, c.into()))
+                .collect(),
+        }
+    }
+}
+
+impl Add<Polynomial<BaseField>> for Polynomial<SecureField> {
+    type Output = Self;
+    fn add(self, rhs: Polynomial<BaseField>) -> Self {
+        self + Polynomial::<SecureField>::from(rhs)
+    }
+}
+
+impl Mul<Polynomial<BaseField>> for Polynomial<SecureField> {
+    type Output = Self;
+    fn mul(self, rhs: Polynomial<BaseField>) -> Self {
+        self * Polynomial::<SecureField>::from(rhs)
+    }
+}
+
 fn to_superscript(n: usize) -> String {
     if n == 0 {
         return std::char::from_u32(0x2070_u32).unwrap().to_string();
@@ -233,20 +306,27 @@ impl Display for Monomial {
 
 impl<F> Display for Polynomial<F>
 where
-    F: Display + Zero + Add + Clone + From<BaseField>,
+    F: Display + Zero + Add + Clone + From<BaseField> + One + PartialEq,
 {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         let mut is_first = true;
+        if self.is_zero() {
+            return write!(f, "0");
+        }
         for (monomial, coef) in self.monomials.iter() {
             if !coef.is_zero() {
                 if !is_first {
                     write!(f, " + ")?;
                 } else {
                     is_first = false;
                 }
-                write!(f, "{}", coef)?;
+                if !coef.is_one() {
+                    write!(f, "{}", coef)?;
+                }
                 if !monomial.vars.is_empty() {
                     write!(f, "{}", monomial)?;
+                } else if coef.is_one() {
+                    write!(f, "1")?;
                 }
             }
         }
@@ -332,6 +412,17 @@ mod tests {
                 vars: [(0, 1), (1, 3), (2, 2)].into()
             }
         );
+
+        let monomial1 = Monomial {
+            vars: [(0, 1), (1, 2)].into(),
+        };
+        let monomial2 = Monomial { vars: [].into() };
+        assert_eq!(
+            monomial1.clone() * monomial2.clone(),
+            Monomial {
+                vars: [(0, 1), (1, 2)].into()
+            }
+        );
     }
 
     #[test]
@@ -346,7 +437,12 @@ mod tests {
             vars: [(4, 1), (5, 2)].into(),
         };
         let poly1 = Polynomial::<M31> {
-            monomials: [(monomial1.clone(), M31(1)), (monomial2.clone(), M31(2))].into(),
+            monomials: [
+                (Monomial::default(), M31(12)),
+                (monomial1.clone(), M31(1)),
+                (monomial2.clone(), M31(2)),
+            ]
+            .into(),
         };
         let poly2 = Polynomial::<M31> {
             monomials: [(monomial2.clone(), M31(5)), (monomial3.clone(), -M31(8))].into(),
@@ -355,6 +451,8 @@ mod tests {
         assert_eq!(
             (poly1.clone() * poly2.clone()).monomials,
             [
+                (monomial2.clone(), M31(60)),
+                (monomial3.clone(), -M31(96)),
                 (monomial1.clone() * monomial2.clone(), M31(5)),
                 (monomial1.clone() * monomial3.clone(), -M31(8)),
                 (monomial2.clone() * monomial2.clone(), M31(10)),

crates/prover/src/constraint_framework/poly_eval.rs

@@ -0,0 +1,83 @@
+use num_traits::One;
+
+use super::rational::Rational;
+use super::EvalAtRow;
+use crate::core::fields::m31::{BaseField, M31};
+use crate::core::fields::qm31::SecureField;
+use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
+
+pub struct PolyEvaluator {
+    pub cur_var_index: usize,
+    pub constraints: Vec<Rational<SecureField>>,
+}
+
+impl EvalAtRow for PolyEvaluator {
+    type F = Rational<BaseField>;
+    type EF = Rational<SecureField>;
+    fn next_interaction_mask<const N: usize>(
+        &mut self,
+        _interaction: usize,
+        offsets: [isize; N],
+    ) -> [Self::F; N] {
+        // TODO(alont) support non-zero offsets.
+        assert_eq!(offsets, [0; N]);
+        self.cur_var_index += 1;
+        std::array::from_fn(|_| Self::F::from_var_index(self.cur_var_index - 1))
+    }
+    fn add_constraint<G>(&mut self, constraint: G)
+    where
+        Self::EF: std::ops::Mul<G, Output = Self::EF>,
+    {
+        self.constraints.push(Self::EF::one() * constraint)
+    }
+    fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF {
+        let [one, i, u, iu] = [
+            SecureField::one(),
+            SecureField::from_m31_array([M31(0), M31(1), M31(0), M31(0)]),
+            SecureField::from_m31_array([M31(0), M31(0), M31(1), M31(0)]),
+            SecureField::from_m31_array([M31(0), M31(0), M31(0), M31(1)]),
+        ];
+        values[0].clone() * one
+            + values[1].clone() * i
+            + values[2].clone() * u
+            + values[3].clone() * iu
+    }
+}
+
+#[cfg(test)]
+mod tests {
+
+    use super::PolyEvaluator;
+    use crate::constraint_framework::{EvalAtRow, FrameworkEval};
+    use crate::core::fields::FieldExpOps;
+    struct TestStruct {}
+    impl FrameworkEval for TestStruct {
+        fn log_size(&self) -> u32 {
+            1 << 16
+        }
+
+        fn max_constraint_log_degree_bound(&self) -> u32 {
+            1 << 17
+        }
+
+        fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
+            let x0 = eval.next_trace_mask();
+            let x1 = eval.next_trace_mask();
+            let x2 = eval.next_trace_mask();
+
+            eval.add_constraint(x0.clone() * x1.clone() * x2 * (x0 + x1).inverse());
+            eval
+        }
+    }
+
+    #[test]
+    fn test_poly_eval() {
+        let test_struct = TestStruct {};
+        let eval = test_struct.evaluate(PolyEvaluator {
+            cur_var_index: 0,
+            constraints: vec![],
+        });
+
+        assert_eq!(eval.constraints[0].to_string(), "(x₀x₁x₂) / (x₀ + x₁)");
+    }
+}