[Move] Compute Lagrange denominator once (#367)

Author: jonas-lj

move/seal/sources/polynomial.move

@@ -4,6 +4,7 @@
 module seal::polynomial;
 
 use seal::gf256;
+use std::u64::range_do_eq;
 
 const EIncomatibleInputLengths: u64 = 1;
 
@@ -15,10 +16,13 @@ public struct Polynomial has copy, drop, store {
 
 /// Evaluate a polynomial at a given point.
 public fun evaluate(p: &Polynomial, x: u8): u8 {
-    let mut result = 0;
+    if (p.coefficients.is_empty()) {
+        return 0
+    };
     let n = p.coefficients.length();
-    n.do!(|i| {
-        result = gf256::add(gf256::mul(result, x), p.coefficients[n - i - 1]);
+    let mut result = p.coefficients[n - 1];
+    (n - 1).do!(|i| {
+        result = gf256::add(gf256::mul(result, x), p.coefficients[n - i - 2]);
     });
     result
 }
@@ -28,39 +32,75 @@ public(package) fun get_constant_term(p: &Polynomial): u8 {
     else p.coefficients[0]
 }
 
-/// Interpolate a polynomial p such that p(x_i) = y[i] for all i.
-/// Panics if the lengths of x and y are not the same.
-/// Panics if x contains duplicate values.
-public(package) fun interpolate(x: &vector<u8>, y: &vector<u8>): Polynomial {
+// Divide a polynomial by the monic linear polynomial x + c.
+// This assumes that the polynomial is divisible by the monic linear polynomial,
+// and it's not clear what the result will be otherwise.
+// Aborts if the polynomial is empty.
+fun div_exact_by_monic_linear(x: &Polynomial, c: u8): Polynomial {
+    let n = x.coefficients.length();
+    let mut coefficients = vector::empty();
+
+    let mut previous = x.coefficients[n - 1];
+    coefficients.push_back(previous);
+
+    range_do_eq!(1, n - 2, |i| {
+        previous = gf256::sub(x.coefficients[n - i - 1], gf256::mul(previous, c));
+        coefficients.push_back(previous);
+    });
+    coefficients.reverse();
+    Polynomial { coefficients }
+}
+
+/// Same as interpolate, but the numerator product, \prod_i (x - x_i), is precomputed.
+fun interpolate_with_numerators(
+    x: &vector<u8>,
+    y: &vector<u8>,
+    numerators: &vector<Polynomial>,
+): Polynomial {
     assert!(x.length() == y.length(), EIncomatibleInputLengths);
     let n = x.length();
     let mut sum = Polynomial { coefficients: vector[] };
     n.do!(|j| {
-        let mut product = Polynomial { coefficients: vector[1] };
+        let mut denominator = 1;
         n.do!(|i| {
             if (i != j) {
-                product =
-                    mul(
-                        &product,
-                        &div(&monic_linear(&x[i]), gf256::sub(x[j], x[i])),
-                    );
+                denominator = gf256::mul(denominator, gf256::sub(x[j], x[i]));
             };
         });
-        sum = add(&sum, &scale(&product, y[j]));
+        sum =
+            add(
+                &sum,
+                &numerators[j].scale(
+                    gf256::div(y[j], denominator),
+                ),
+            );
     });
     sum
 }
 
+/// Compute the numerators of the Lagrange polynomials for the given x values.
+fun compute_numerators(x: vector<u8>): vector<Polynomial> {
+    // The full numerator depends only on x, so we can compute it here
+    let full_numerator = x.fold!(Polynomial { coefficients: vector[1] }, |product, x_j| {
+        product.mul(&monic_linear(&x_j))
+    });
+    x.map_ref!(|x_j| div_exact_by_monic_linear(&full_numerator, *x_j))
+}
+
 /// Interpolate l polynomials p_1, ..., p_l such that p_i(x_j) = y[j][i] for all i, j.
 /// The length of the input vectors must be the same.
 /// The length of each vector in y must be the same (equal to the l above).
+/// Aborts if the input lengths are not compatible or if the vectors are empty.
 public(package) fun interpolate_all(x: &vector<u8>, y: &vector<vector<u8>>): vector<Polynomial> {
     assert!(x.length() == y.length(), EIncomatibleInputLengths);
     let l = y[0].length();
     assert!(y.all!(|yi| yi.length() == l), EIncomatibleInputLengths);
+
+    // The numerators depend only on x, so we can compute them here
+    let numerators = compute_numerators(*x);
     vector::tabulate!(l, |i| {
         let yi = y.map_ref!(|yj| yj[i]);
-        interpolate(x, &yi)
+        interpolate_with_numerators(x, &yi, &numerators)
     })
 }
 
@@ -101,17 +141,13 @@ fun mul(x: &Polynomial, y: &Polynomial): Polynomial {
     Polynomial { coefficients }
 }
 
-fun div(x: &Polynomial, s: u8): Polynomial {
-    x.scale(gf256::div(1, s))
-}
-
 fun scale(x: &Polynomial, s: u8): Polynomial {
     Polynomial { coefficients: x.coefficients.map_ref!(|c| gf256::mul(*c, s)) }
 }
 
-/// Return x - c
+/// Return x - c (same as x + c since GF256 is a binary field)
 fun monic_linear(c: &u8): Polynomial {
-    Polynomial { coefficients: vector[gf256::sub(0, *c), 1] }
+    Polynomial { coefficients: vector[*c, 1] }
 }
 
 #[test]
@@ -133,13 +169,23 @@ fun test_evaluate() {
     assert!(p.evaluate(1) == 0);
     assert!(p.evaluate(2) == 9);
     assert!(p.evaluate(3) == 8);
+
+    // Test zero polynomial
+    let p = Polynomial { coefficients: vector[] };
+    assert!(p.evaluate(0) == 0);
+
+    let p = Polynomial { coefficients: vector[3] };
+    assert!(p.evaluate(0) == 3);
+    assert!(p.evaluate(1) == 3);
+    assert!(p.evaluate(2) == 3);
+    assert!(p.evaluate(3) == 3);
 }
 
 #[test]
 fun test_interpolate() {
     let x = vector[1, 2, 3];
     let y = vector[7, 11, 17];
-    let p = interpolate(&x, &y);
+    let p = interpolate_with_numerators(&x, &y, &compute_numerators(x));
     assert!(p.coefficients == x"1d150f");
     x.zip_do!(y, |x, y| assert!(p.evaluate(x) == y));
 }
@@ -155,3 +201,12 @@ fun test_interpolate_all() {
         assert!(ps[1].evaluate(x) == y[1]);
     });
 }
+
+#[test]
+fun test_div_exact_by_monic_linear() {
+    let x = Polynomial { coefficients: vector[1, 2, 3, 4, 5, 6, 7] };
+    let monic_linear = monic_linear(&2);
+    let y = mul(&x, &monic_linear);
+    let z = div_exact_by_monic_linear(&y, 2);
+    assert!(z == x);
+}