From 7651e39f7a1e8029d7c33ecb3ff7a0222e4057fc Mon Sep 17 00:00:00 2001
From: zhenfei <zhenfei.zhang@hotmail.com>
Date: Mon, 16 May 2022 20:40:46 -0400
Subject: [PATCH] optimize uni poly interpolation (#19)

---
 poly-iop/Cargo.toml                |  4 +-
 poly-iop/src/sum_check/verifier.rs | 86 ++++++++++++++++++++++++------
 poly-iop/src/virtual_poly.rs       |  6 +--
 3 files changed, 74 insertions(+), 22 deletions(-)

diff --git a/poly-iop/Cargo.toml b/poly-iop/Cargo.toml
index 3fcf4bd..1ee3835 100644
--- a/poly-iop/Cargo.toml
+++ b/poly-iop/Cargo.toml
@@ -26,8 +26,8 @@ path = "benches/bench.rs"
 harness = false
 
 [features]
-# default = [ "parallel", "print-trace" ]
-default = [ "parallel" ]
+default = [ "parallel", "print-trace" ]
+# default = [ "parallel" ]
 parallel = [ 
     "rayon",
     "ark-std/parallel", 
diff --git a/poly-iop/src/sum_check/verifier.rs b/poly-iop/src/sum_check/verifier.rs
index 7a0b499..8a3a4f6 100644
--- a/poly-iop/src/sum_check/verifier.rs
+++ b/poly-iop/src/sum_check/verifier.rs
@@ -116,7 +116,7 @@ impl<F: PrimeField> SumCheckVerifier<F> for IOPVerifierState<F> {
                         self.max_degree + 1
                     )));
                 }
-                Ok(interpolate_uni_poly::<F>(&evaluations, challenge))
+                interpolate_uni_poly::<F>(&evaluations, challenge)
             })
             .collect::<Result<Vec<_>, PolyIOPErrors>>()?;
 
@@ -134,7 +134,7 @@ impl<F: PrimeField> SumCheckVerifier<F> for IOPVerifierState<F> {
                         self.max_degree + 1
                     )));
                 }
-                Ok(interpolate_uni_poly::<F>(&evaluations, challenge))
+                interpolate_uni_poly::<F>(&evaluations, challenge)
             })
             .collect::<Result<Vec<_>, PolyIOPErrors>>()?;
         // insert the asserted_sum to the first position of the expected vector
@@ -162,23 +162,75 @@ impl<F: PrimeField> SumCheckVerifier<F> for IOPVerifierState<F> {
 }
 
 /// Interpolate a uni-variate degree-`p_i.len()-1` polynomial and evaluate this
-/// polynomial at `eval_at`.
-pub(crate) fn interpolate_uni_poly<F: PrimeField>(p_i: &[F], eval_at: F) -> F {
-    let start = start_timer!(|| "sum check interpolate uni poly");
-    let mut result = F::zero();
-    let mut i = F::zero();
-    for term in p_i.iter() {
-        let mut term = *term;
-        let mut j = F::zero();
-        for _ in 0..p_i.len() {
+/// polynomial at `eval_at`:
+///   \sum_{i=0}^len p_i * (\prod_{j!=i} (eval_at - j)/(i-j) )
+/// This implementation is linear in number of inputs in terms of field
+/// operations. It also has a quadratic term in primitive operations which is
+/// negligible compared to field operations.
+pub(crate) fn interpolate_uni_poly<F: PrimeField>(
+    p_i: &[F],
+    eval_at: F,
+) -> Result<F, PolyIOPErrors> {
+    let start = start_timer!(|| "sum check interpolate uni poly opt");
+
+    let mut res = F::zero();
+
+    // prod = \prod_{j!=i} (eval_at - j)
+    let mut evals = vec![];
+    let len = p_i.len();
+    let mut prod = eval_at;
+    evals.push(eval_at);
+
+    for e in 1..len {
+        let tmp = eval_at - F::from(e as u64);
+        evals.push(tmp);
+        prod *= tmp;
+    }
+
+    for i in 0..len {
+        let divisor = get_divisor(i, len)?;
+        let divisor_f = {
+            if divisor < 0 {
+                -F::from((-divisor) as u128)
+            } else {
+                F::from(divisor as u128)
+            }
+        };
+        res += p_i[i] * prod / (divisor_f * evals[i]);
+    }
+
+    end_timer!(start);
+    Ok(res)
+}
+
+/// Compute \prod_{j!=i)^len (i-j). This function takes O(n^2) number of
+/// primitive operations which is negligible compared to field operations.
+// We know
+//  - factorial(20) ~ 2^61
+//  - factorial(33) ~ 2^123
+// so we will be able to store the result for len<=20 with i64;
+// for len<=33 with i128; and we do not currently support len>33.
+#[inline]
+fn get_divisor(i: usize, len: usize) -> Result<i128, PolyIOPErrors> {
+    if len <= 20 {
+        let mut res = 1i64;
+        for j in 0..len {
+            if j != i {
+                res *= i as i64 - j as i64;
+            }
+        }
+        Ok(res as i128)
+    } else if len <= 33 {
+        let mut res = 1i128;
+        for j in 0..len {
             if j != i {
-                term = term * (eval_at - j) / (i - j)
+                res *= i as i128 - j as i128;
             }
-            j += F::one();
         }
-        i += F::one();
-        result += term;
+        Ok(res)
+    } else {
+        Err(PolyIOPErrors::InvalidParameters(
+            "Do not support number variable > 33".to_string(),
+        ))
     }
-    end_timer!(start);
-    result
 }
diff --git a/poly-iop/src/virtual_poly.rs b/poly-iop/src/virtual_poly.rs
index 88cf66a..dc27059 100644
--- a/poly-iop/src/virtual_poly.rs
+++ b/poly-iop/src/virtual_poly.rs
@@ -205,7 +205,7 @@ impl<F: PrimeField> VirtualPolynomial<F> {
         num_products: usize,
         rng: &mut R,
     ) -> Result<(Self, F), PolyIOPErrors> {
-        let start = start_timer!("sample random virtual polynomial");
+        let start = start_timer!(|| "sample random virtual polynomial");
 
         let mut sum = F::zero();
         let mut poly = VirtualPolynomial::new(nv);
@@ -252,7 +252,7 @@ fn random_mle_list<F: PrimeField, R: RngCore>(
     degree: usize,
     rng: &mut R,
 ) -> (Vec<Rc<DenseMultilinearExtension<F>>>, F) {
-    let start = start_timer!("sample random mle list");
+    let start = start_timer!(|| "sample random mle list");
     let mut multiplicands = Vec::with_capacity(degree);
     for _ in 0..degree {
         multiplicands.push(Vec::with_capacity(1 << nv))
@@ -285,7 +285,7 @@ pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
     degree: usize,
     rng: &mut R,
 ) -> Vec<Rc<DenseMultilinearExtension<F>>> {
-    let start = start_timer!("sample random zero mle list");
+    let start = start_timer!(|| "sample random zero mle list");
 
     let mut multiplicands = Vec::with_capacity(degree);
     for _ in 0..degree {