Merge pull request #15 from mir-protocol/lagrange_interp

Interpolants of arbitrary (point, value) lists

src/field/crandall_field.rs

@@ -6,6 +6,7 @@ use num::Integer;
 
 use crate::field::field::Field;
 use std::hash::{Hash, Hasher};
+use std::iter::{Product, Sum};
 
 /// EPSILON = 9 * 2**28 - 1
 const EPSILON: u64 = 2415919103;
@@ -257,6 +258,12 @@ impl AddAssign for CrandallField {
     }
 }
 
+impl Sum for CrandallField {
+    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
+        iter.fold(Self::ZERO, |acc, x| acc + x)
+    }
+}
+
 impl Sub for CrandallField {
     type Output = Self;
 
@@ -291,6 +298,12 @@ impl MulAssign for CrandallField {
     }
 }
 
+impl Product for CrandallField {
+    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
+        iter.fold(Self::ONE, |acc, x| acc * x)
+    }
+}
+
 impl Div for CrandallField {
     type Output = Self;
 

src/field/fft.rs

@@ -169,7 +169,7 @@ mod tests {
         for i in 0..degree {
             coefficients.push(F::from_canonical_usize(i * 1337 % 100));
         }
-        let coefficients = PolynomialCoeffs::pad(coefficients);
+        let coefficients = PolynomialCoeffs::new_padded(coefficients);
 
         let points = fft(coefficients.clone());
         assert_eq!(points, evaluate_naive(&coefficients));

src/field/field.rs

@@ -1,10 +1,13 @@
-use crate::util::bits_u64;
-use rand::rngs::OsRng;
-use rand::Rng;
 use std::fmt::{Debug, Display};
 use std::hash::Hash;
+use std::iter::{Product, Sum};
 use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
 
+use rand::rngs::OsRng;
+use rand::Rng;
+
+use crate::util::bits_u64;
+
 /// A finite field with prime order less than 2^64.
 pub trait Field:
     'static
@@ -14,10 +17,12 @@ pub trait Field:
     + Neg<Output = Self>
     + Add<Self, Output = Self>
     + AddAssign<Self>
+    + Sum
     + Sub<Self, Output = Self>
     + SubAssign<Self>
     + Mul<Self, Output = Self>
     + MulAssign<Self>
+    + Product
     + Div<Self, Output = Self>
     + DivAssign<Self>
     + Debug
@@ -42,6 +47,10 @@ pub trait Field:
         *self == Self::ZERO
     }
 
+    fn is_nonzero(&self) -> bool {
+        *self != Self::ZERO
+    }
+
     fn is_one(&self) -> bool {
         *self == Self::ONE
     }
@@ -90,12 +99,12 @@ pub trait Field:
         x_inv
     }
 
-    fn primitive_root_of_unity(n_power: usize) -> Self {
+    fn primitive_root_of_unity(n_log: usize) -> Self {
-        assert!(n_power <= Self::TWO_ADICITY);
+        assert!(n_log <= Self::TWO_ADICITY);
         let base = Self::POWER_OF_TWO_GENERATOR;
         // TODO: Just repeated squaring should be a bit faster, to avoid conditionals.
         base.exp(Self::from_canonical_u64(
-            1u64 << (Self::TWO_ADICITY - n_power),
+            1u64 << (Self::TWO_ADICITY - n_log),
         ))
     }
 

src/field/lagrange.rs

@@ -0,0 +1,132 @@
+use crate::field::fft::ifft;
+use crate::field::field::Field;
+use crate::polynomial::polynomial::{PolynomialCoeffs, PolynomialValues};
+use crate::util::log2_ceil;
+
+/// Computes the unique degree < n interpolant of an arbitrary list of n (point, value) pairs.
+///
+/// Note that the implementation assumes that `F` is two-adic, in particular that
+/// `2^{F::TWO_ADICITY} >= points.len()`. This leads to a simple FFT-based implementation.
+pub(crate) fn interpolant<F: Field>(points: &[(F, F)]) -> PolynomialCoeffs<F> {
+    let n = points.len();
+    let n_log = log2_ceil(n);
+    let n_padded = 1 << n_log;
+
+    let g = F::primitive_root_of_unity(n_log);
+    let subgroup = F::cyclic_subgroup_known_order(g, n_padded);
+    let barycentric_weights = barycentric_weights(points);
+    let subgroup_evals = subgroup
+        .into_iter()
+        .map(|x| interpolate(points, x, &barycentric_weights))
+        .collect();
+
+    let mut coeffs = ifft(PolynomialValues {
+        values: subgroup_evals,
+    });
+    coeffs.trim();
+    coeffs
+}
+
+/// Interpolate the polynomial defined by an arbitrary set of (point, value) pairs at the given
+/// point `x`.
+fn interpolate<F: Field>(points: &[(F, F)], x: F, barycentric_weights: &[F]) -> F {
+    // If x is in the list of points, the Lagrange formula would divide by zero.
+    for &(x_i, y_i) in points {
+        if x_i == x {
+            return y_i;
+        }
+    }
+
+    let l_x: F = points.iter().map(|&(x_i, y_i)| x - x_i).product();
+
+    let sum = (0..points.len())
+        .map(|i| {
+            let x_i = points[i].0;
+            let y_i = points[i].1;
+            let w_i = barycentric_weights[i];
+            w_i / (x - x_i) * y_i
+        })
+        .sum();
+
+    l_x * sum
+}
+
+fn barycentric_weights<F: Field>(points: &[(F, F)]) -> Vec<F> {
+    let n = points.len();
+    (0..n)
+        .map(|i| {
+            (0..n)
+                .filter(|&j| j != i)
+                .map(|j| points[i].0 - points[j].0)
+                .product::<F>()
+                .inverse()
+        })
+        .collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::field::crandall_field::CrandallField;
+    use crate::field::field::Field;
+    use crate::field::lagrange::interpolant;
+    use crate::polynomial::polynomial::PolynomialCoeffs;
+
+    #[test]
+    fn interpolant_random() {
+        type F = CrandallField;
+
+        for deg in 0..10 {
+            let domain = (0..deg).map(|_| F::rand()).collect::<Vec<_>>();
+            let coeffs = (0..deg).map(|_| F::rand()).collect();
+            let coeffs = PolynomialCoeffs { coeffs };
+
+            let points = eval_naive(&coeffs, &domain);
+            assert_eq!(interpolant(&points), coeffs);
+        }
+    }
+
+    #[test]
+    fn interpolant_random_roots_of_unity() {
+        type F = CrandallField;
+
+        for deg_log in 0..4 {
+            let deg = 1 << deg_log;
+            let g = F::primitive_root_of_unity(deg_log);
+            let domain = F::cyclic_subgroup_known_order(g, deg);
+            let coeffs = (0..deg).map(|_| F::rand()).collect();
+            let coeffs = PolynomialCoeffs { coeffs };
+
+            let points = eval_naive(&coeffs, &domain);
+            assert_eq!(interpolant(&points), coeffs);
+        }
+    }
+
+    #[test]
+    fn interpolant_random_overspecified() {
+        type F = CrandallField;
+
+        for deg in 0..10 {
+            let points = deg + 5;
+            let domain = (0..points).map(|_| F::rand()).collect::<Vec<_>>();
+            let coeffs = (0..deg).map(|_| F::rand()).collect();
+            let coeffs = PolynomialCoeffs { coeffs };
+
+            let points = eval_naive(&coeffs, &domain);
+            assert_eq!(interpolant(&points), coeffs);
+        }
+    }
+
+    fn eval_naive<F: Field>(coeffs: &PolynomialCoeffs<F>, domain: &[F]) -> Vec<(F, F)> {
+        domain
+            .iter()
+            .map(|&x| {
+                let eval = x
+                    .powers()
+                    .zip(&coeffs.coeffs)
+                    .map(|(x_power, &coeff)| coeff * x_power)
+                    .sum();
+                (x, eval)
+            })
+            .collect()
+    }
+}

src/field/mod.rs

@@ -2,6 +2,7 @@ pub(crate) mod cosets;
 pub mod crandall_field;
 pub mod fft;
 pub mod field;
+pub(crate) mod lagrange;
 
 #[cfg(test)]
 mod field_testing;

src/polynomial/polynomial.rs

@@ -2,7 +2,7 @@ use crate::field::fft::{fft, ifft};
 use crate::field::field::Field;
 use crate::util::log2_strict;
 
-/// A polynomial in point-value form. The number of values must be a power of two.
+/// A polynomial in point-value form.
 ///
 /// The points are implicitly `g^i`, where `g` generates the subgroup whose size equals the number
 /// of points.
@@ -13,7 +13,6 @@ pub struct PolynomialValues<F: Field> {
 
 impl<F: Field> PolynomialValues<F> {
     pub fn new(values: Vec<F>) -> Self {
-        assert!(values.len().is_power_of_two());
         PolynomialValues { values }
     }
 
@@ -36,7 +35,7 @@ impl<F: Field> PolynomialValues<F> {
     }
 }
 
-/// A polynomial in coefficient form. The number of coefficients must be a power of two.
+/// A polynomial in coefficient form.
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct PolynomialCoeffs<F: Field> {
     pub(crate) coeffs: Vec<F>,
@@ -44,11 +43,11 @@ pub struct PolynomialCoeffs<F: Field> {
 
 impl<F: Field> PolynomialCoeffs<F> {
     pub fn new(coeffs: Vec<F>) -> Self {
-        assert!(coeffs.len().is_power_of_two());
         PolynomialCoeffs { coeffs }
     }
 
-    pub(crate) fn pad(mut coeffs: Vec<F>) -> Self {
+    /// Create a new polynomial with its coefficient list padded to the next power of two.
+    pub(crate) fn new_padded(mut coeffs: Vec<F>) -> Self {
         while !coeffs.len().is_power_of_two() {
             coeffs.push(F::ZERO);
         }
@@ -70,7 +69,6 @@ impl<F: Field> PolynomialCoeffs<F> {
     }
 
     pub(crate) fn chunks(&self, chunk_size: usize) -> Vec<Self> {
-        assert!(chunk_size.is_power_of_two());
         self.coeffs
             .chunks(chunk_size)
             .map(|chunk| PolynomialCoeffs::new(chunk.to_vec()))
@@ -97,4 +95,17 @@ impl<F: Field> PolynomialCoeffs<F> {
         }
         Self { coeffs }
     }
+
+    /// Removes leading zero coefficients.
+    pub fn trim(&mut self) {
+        self.coeffs.drain(self.degree_plus_one()..);
+    }
+
+    /// Degree of the polynomial + 1.
+    fn degree_plus_one(&self) -> usize {
+        (0usize..self.len())
+            .rev()
+            .find(|&i| self.coeffs[i].is_nonzero())
+            .map_or(0, |i| i + 1)
+    }
 }