From 3ea5b61c5e713282cb983b52804a0c83f549cbd7 Mon Sep 17 00:00:00 2001
From: zhenfei <zhenfei.zhang@hotmail.com>
Date: Fri, 20 May 2022 16:29:14 -0400
Subject: [PATCH] 26 implement `prod(0,x)` (#27)

---
 poly-iop/src/lib.rs            |  1 +
 poly-iop/src/perm_check/mod.rs | 92 ++++++++++++++++++++++++++++++++++
 poly-iop/src/utils.rs          | 11 ++++
 poly-iop/src/virtual_poly.rs   | 13 +----
 4 files changed, 106 insertions(+), 11 deletions(-)
 create mode 100644 poly-iop/src/perm_check/mod.rs

diff --git a/poly-iop/src/lib.rs b/poly-iop/src/lib.rs
index 8e2fff0..db9dbf2 100644
--- a/poly-iop/src/lib.rs
+++ b/poly-iop/src/lib.rs
@@ -2,6 +2,7 @@ use ark_ff::PrimeField;
 use std::marker::PhantomData;
 
 mod errors;
+mod perm_check;
 mod structs;
 mod sum_check;
 mod transcript;
diff --git a/poly-iop/src/perm_check/mod.rs b/poly-iop/src/perm_check/mod.rs
new file mode 100644
index 0000000..9de2eda
--- /dev/null
+++ b/poly-iop/src/perm_check/mod.rs
@@ -0,0 +1,92 @@
+//! This module implements the permutation check protocol.
+
+use ark_ff::PrimeField;
+use ark_poly::DenseMultilinearExtension;
+
+use crate::PolyIOPErrors;
+
+/// Compute `prod(0,x) := prod(0, x1, …, xn)` which is the MLE over the
+/// evaluations of the following polynomial on the boolean hypercube {0,1}^n:
+///
+///  (w(x) + \beta s_id(x) + \gamma)/(w(x) + \beta s_perm(x) + \gamma)
+///
+/// where
+/// - beta and gamma are challenges
+/// - w(x), s_id(x), s_perm(x) are mle-s
+#[allow(dead_code)]
+// TODO: remove
+fn compute_prod_0<F: PrimeField>(
+    beta: &F,
+    gamma: &F,
+    w: &DenseMultilinearExtension<F>,
+    s_id: &DenseMultilinearExtension<F>,
+    s_perm: &DenseMultilinearExtension<F>,
+) -> Result<DenseMultilinearExtension<F>, PolyIOPErrors> {
+    let num_vars = w.num_vars;
+
+    if num_vars != s_id.num_vars || num_vars != s_perm.num_vars {
+        return Err(PolyIOPErrors::InvalidParameters(
+            "num of variables do not match".to_string(),
+        ));
+    }
+    let eval: Vec<F> = w
+        .iter()
+        .zip(s_id.iter().zip(s_perm.iter()))
+        .map(|(wi, (s_id_i, s_perm_i))| {
+            let tmp = *wi + *gamma;
+            (tmp + *beta * *s_id_i) / (tmp + *beta * *s_perm_i)
+        })
+        .collect();
+
+    let res = DenseMultilinearExtension::from_evaluations_vec(num_vars, eval);
+
+    Ok(res)
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::utils::bit_decompose;
+    use ark_bls12_381::Fr;
+    use ark_ff::UniformRand;
+    use ark_poly::MultilinearExtension;
+    use ark_std::test_rng;
+
+    #[test]
+    fn test_compute_prod_0() -> Result<(), PolyIOPErrors> {
+        let mut rng = test_rng();
+
+        for num_vars in 2..6 {
+            let w_vec: Vec<Fr> = (0..(1 << num_vars)).map(|_| Fr::rand(&mut rng)).collect();
+            let w = DenseMultilinearExtension::from_evaluations_vec(num_vars, w_vec);
+
+            let s_id_vec: Vec<Fr> = (0..(1 << num_vars)).map(|_| Fr::rand(&mut rng)).collect();
+            let s_id = DenseMultilinearExtension::from_evaluations_vec(num_vars, s_id_vec);
+
+            let s_perm_vec: Vec<Fr> = (0..(1 << num_vars)).map(|_| Fr::rand(&mut rng)).collect();
+            let s_perm = DenseMultilinearExtension::from_evaluations_vec(num_vars, s_perm_vec);
+
+            let beta = Fr::rand(&mut rng);
+            let gamma = Fr::rand(&mut rng);
+            for i in 0..1 << num_vars {
+                let r: Vec<Fr> = bit_decompose(i, num_vars)
+                    .iter()
+                    .map(|&x| Fr::from(x))
+                    .collect();
+
+                let prod_0 = compute_prod_0(&beta, &gamma, &w, &s_id, &s_perm)?;
+
+                let eval = prod_0.evaluate(&r).unwrap();
+
+                let w_eval = w.evaluate(&r).unwrap();
+                let s_id_eval = s_id.evaluate(&r).unwrap();
+                let s_perm_eval = s_perm.evaluate(&r).unwrap();
+                let eval_rec =
+                    (w_eval + beta * s_id_eval + gamma) / (w_eval + beta * s_perm_eval + gamma);
+
+                assert_eq!(eval, eval_rec);
+            }
+        }
+        Ok(())
+    }
+}
diff --git a/poly-iop/src/utils.rs b/poly-iop/src/utils.rs
index 831c2fe..c11d665 100644
--- a/poly-iop/src/utils.rs
+++ b/poly-iop/src/utils.rs
@@ -11,6 +11,17 @@ macro_rules! to_bytes {
     }};
 }
 
+#[cfg(test)]
+pub(crate) fn bit_decompose(input: u64, num_var: usize) -> Vec<bool> {
+    let mut res = Vec::with_capacity(num_var);
+    let mut i = input;
+    for _ in 0..num_var {
+        res.push(i & 1 == 1);
+        i >>= 1;
+    }
+    res
+}
+
 #[test]
 fn test_to_bytes() {
     use ark_bls12_381::Fr;
diff --git a/poly-iop/src/virtual_poly.rs b/poly-iop/src/virtual_poly.rs
index 49d4297..3473a1c 100644
--- a/poly-iop/src/virtual_poly.rs
+++ b/poly-iop/src/virtual_poly.rs
@@ -436,8 +436,9 @@ fn build_eq_x_r_helper<F: PrimeField>(r: &[F], buf: &mut Vec<F>) -> Result<(), P
 }
 
 #[cfg(test)]
-pub(crate) mod test {
+mod test {
     use super::*;
+    use crate::utils::bit_decompose;
     use ark_bls12_381::Fr;
     use ark_ff::UniformRand;
     use ark_std::test_rng;
@@ -548,14 +549,4 @@ pub(crate) mod test {
         end_timer!(start);
         res
     }
-
-    fn bit_decompose(input: u64, num_var: usize) -> Vec<bool> {
-        let mut res = Vec::with_capacity(num_var);
-        let mut i = input;
-        for _ in 0..num_var {
-            res.push(i & 1 == 1);
-            i >>= 1;
-        }
-        res
-    }
 }