Added rational functions for formal polynomial evaluator.

crates/prover/src/constraint_framework/mod.rs

@@ -8,6 +8,7 @@ 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,13 +1,12 @@
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
-use std::ops::{Add, AddAssign, Mul, MulAssign, Neg};
+use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub};
 
 use itertools::Itertools;
 use num_traits::{One, Zero};
 
 use crate::core::fields::m31::BaseField;
 use crate::core::fields::qm31::SecureField;
-use crate::core::fields::FieldExpOps;
 
 /// A monic monomial consists of a list of variables and their exponents.
 #[derive(Debug, Hash, PartialEq, PartialOrd, Eq, Ord, Clone)]
@@ -22,9 +21,13 @@ impl Monomial {
     }
 
     fn default() -> Monomial {
-        Monomial {
-            vars: [(0, 0)].into(),
-        }
+        Monomial { vars: [].into() }
+    }
+}
+
+impl One for Monomial {
+    fn one() -> Self {
+        Monomial { vars: [].into() }
     }
 }
 
@@ -34,6 +37,21 @@ 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 {
@@ -73,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;
@@ -199,14 +227,29 @@ where
     }
 }
 
-impl<F> FieldExpOps for Polynomial<F>
-where
-    F: Zero + One + Clone + Add + Mul + AddAssign + From<BaseField>,
-{
-    fn inverse(&self) -> Self {
-        assert!(!self.is_zero(), "0 has no inverse");
-        let mut res = Self::from(BaseField::one()) / self.clone();
-        res
+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)
     }
 }
 
@@ -263,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")?;
                 }
             }
         }
@@ -362,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]
@@ -376,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(),
@@ -385,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

@@ -1,28 +1,83 @@
-use super::poly::Polynomial;
+use num_traits::One;
+
+use super::rational::Rational;
 use super::EvalAtRow;
-use crate::core::fields::m31::BaseField;
+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<'a> {}
+pub struct PolyEvaluator {
+    pub cur_var_index: usize,
+    pub constraints: Vec<Rational<SecureField>>,
+}
 
-impl<'a> EvalAtRow for PolyEvaluator<'a> {
-    type F = Polynomial<BaseField>;
-    type EF = Polynomial<SecureField>;
+impl EvalAtRow for PolyEvaluator {
+    type F = Rational<BaseField>;
+    type EF = Rational<SecureField>;
     fn next_interaction_mask<const N: usize>(
         &mut self,
-        interaction: usize,
+        _interaction: usize,
         offsets: [isize; N],
     ) -> [Self::F; N] {
-        unimplemented!()
+        // 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>,
     {
-        unimplemented!()
+        self.constraints.push(Self::EF::one() * constraint)
     }
     fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF {
-        unimplemented!()
+        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₁)");
     }
 }