enabling batch opening and mock tests (#80)

- add mock circuits
- add vanilla and jellyfish plonk gates
- performance tuning

arithmetic/Cargo.toml

@@ -19,6 +19,7 @@ rayon = { version = "1.5.2", default-features = false, optional = true }
 
 [dev-dependencies]
 ark-ec = { version = "^0.3.0", default-features = false }
+criterion = "0.3.0"
 
 [features]
 # default = [ "parallel", "print-trace" ]
@@ -31,4 +32,10 @@ parallel = [
     ]
 print-trace = [ 
     "ark-std/print-trace" 
-    ] 
+    ] 
+
+
+[[bench]]
+name = "mle_eval"
+path = "benches/bench.rs"
+harness = false

arithmetic/benches/bench.rs

@@ -0,0 +1,37 @@
+#[macro_use]
+extern crate criterion;
+
+use arithmetic::fix_variables;
+use ark_bls12_381::Fr;
+use ark_ff::Field;
+use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
+use ark_std::{ops::Range, test_rng};
+use criterion::{black_box, BenchmarkId, Criterion};
+
+const NUM_VARIABLES_RANGE: Range<usize> = 10..21;
+
+fn evaluation_op_bench<F: Field>(c: &mut Criterion) {
+    let mut rng = test_rng();
+    let mut group = c.benchmark_group("Evaluate");
+    for nv in NUM_VARIABLES_RANGE {
+        group.bench_with_input(BenchmarkId::new("evaluate native", nv), &nv, |b, &nv| {
+            let poly = DenseMultilinearExtension::<F>::rand(nv, &mut rng);
+            let point: Vec<_> = (0..nv).map(|_| F::rand(&mut rng)).collect();
+            b.iter(|| black_box(poly.evaluate(&point).unwrap()))
+        });
+
+        group.bench_with_input(BenchmarkId::new("evaluate optimized", nv), &nv, |b, &nv| {
+            let poly = DenseMultilinearExtension::<F>::rand(nv, &mut rng);
+            let point: Vec<_> = (0..nv).map(|_| F::rand(&mut rng)).collect();
+            b.iter(|| black_box(fix_variables(&poly, &point)))
+        });
+    }
+    group.finish();
+}
+
+fn bench_bls_381(c: &mut Criterion) {
+    evaluation_op_bench::<Fr>(c);
+}
+
+criterion_group!(benches, bench_bls_381);
+criterion_main!(benches);

arithmetic/src/lib.rs

@@ -1,7 +1,15 @@
 mod errors;
 mod multilinear_polynomial;
+mod univariate_polynomial;
+mod util;
 mod virtual_polynomial;
 
 pub use errors::ArithErrors;
-pub use multilinear_polynomial::{random_zero_mle_list, DenseMultilinearExtension};
+pub use multilinear_polynomial::{
+    evaluate_no_par, evaluate_opt, fix_first_variable, fix_variables, identity_permutation_mle,
+    merge_polynomials, random_mle_list, random_permutation_mle, random_zero_mle_list,
+    DenseMultilinearExtension,
+};
+pub use univariate_polynomial::{build_l, get_uni_domain};
+pub use util::{bit_decompose, gen_eval_point, get_batched_nv, get_index};
 pub use virtual_polynomial::{build_eq_x_r, VPAuxInfo, VirtualPolynomial};

arithmetic/src/multilinear_polynomial.rs

@@ -1,9 +1,49 @@
-use ark_ff::PrimeField;
+use crate::{util::get_batched_nv, ArithErrors};
+use ark_ff::{Field, PrimeField};
+use ark_poly::MultilinearExtension;
 use ark_std::{end_timer, rand::RngCore, start_timer};
+#[cfg(feature = "parallel")]
+use rayon::prelude::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
 use std::rc::Rc;
 
 pub use ark_poly::DenseMultilinearExtension;
 
+/// Sample a random list of multilinear polynomials.
+/// Returns
+/// - the list of polynomials,
+/// - its sum of polynomial evaluations over the boolean hypercube.
+pub fn random_mle_list<F: PrimeField, R: RngCore>(
+    nv: usize,
+    degree: usize,
+    rng: &mut R,
+) -> (Vec<Rc<DenseMultilinearExtension<F>>>, F) {
+    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))
+    }
+    let mut sum = F::zero();
+
+    for _ in 0..(1 << nv) {
+        let mut product = F::one();
+
+        for e in multiplicands.iter_mut() {
+            let val = F::rand(rng);
+            e.push(val);
+            product *= val;
+        }
+        sum += product;
+    }
+
+    let list = multiplicands
+        .into_iter()
+        .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
+        .collect();
+
+    end_timer!(start);
+    (list, sum)
+}
+
 // Build a randomize list of mle-s whose sum is zero.
 pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
     nv: usize,
@@ -31,3 +71,142 @@ pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
     end_timer!(start);
     list
 }
+
+/// An MLE that represent an identity permutation: `f(index) \mapto index`
+pub fn identity_permutation_mle<F: PrimeField>(
+    num_vars: usize,
+) -> Rc<DenseMultilinearExtension<F>> {
+    let s_id_vec = (0..1u64 << num_vars).map(F::from).collect();
+    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, s_id_vec,
+    ))
+}
+
+/// An MLE that represent a random permutation
+pub fn random_permutation_mle<F: PrimeField, R: RngCore>(
+    num_vars: usize,
+    rng: &mut R,
+) -> Rc<DenseMultilinearExtension<F>> {
+    let len = 1u64 << num_vars;
+    let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
+    let mut s_perm_vec = vec![];
+    for _ in 0..len {
+        let index = rng.next_u64() as usize % s_id_vec.len();
+        s_perm_vec.push(s_id_vec.remove(index));
+    }
+    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, s_perm_vec,
+    ))
+}
+
+pub fn evaluate_opt<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
+    assert_eq!(poly.num_vars, point.len());
+    fix_variables(poly, point).evaluations[0]
+}
+
+pub fn fix_variables<F: Field>(
+    poly: &DenseMultilinearExtension<F>,
+    partial_point: &[F],
+) -> DenseMultilinearExtension<F> {
+    assert!(
+        partial_point.len() <= poly.num_vars,
+        "invalid size of partial point"
+    );
+    let nv = poly.num_vars;
+    let mut poly = poly.evaluations.to_vec();
+    let dim = partial_point.len();
+    // evaluate single variable of partial point from left to right
+    for (i, point) in partial_point.iter().enumerate().take(dim) {
+        poly = fix_one_variable_helper(&poly, nv - i, point);
+    }
+
+    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
+}
+
+pub fn fix_first_variable<F: Field>(
+    poly: &DenseMultilinearExtension<F>,
+    partial_point: &F,
+) -> DenseMultilinearExtension<F> {
+    assert!(poly.num_vars != 0, "invalid size of partial point");
+
+    let nv = poly.num_vars;
+    let res = fix_one_variable_helper(&poly.evaluations, nv, partial_point);
+    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - 1, &res)
+}
+
+fn fix_one_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
+    let mut res = vec![F::zero(); 1 << (nv - 1)];
+    let one_minus_p = F::one() - point;
+
+    // evaluate single variable of partial point from left to right
+    #[cfg(not(feature = "parallel"))]
+    for b in 0..(1 << (nv - 1)) {
+        res[b] = data[b << 1] * one_minus_p + data[(b << 1) + 1] * point;
+    }
+
+    #[cfg(feature = "parallel")]
+    if nv >= 13 {
+        // on my computer we parallelization doesn't help till nv >= 13
+        res.par_iter_mut().enumerate().for_each(|(i, x)| {
+            *x = data[i << 1] * one_minus_p + data[(i << 1) + 1] * point;
+        });
+    } else {
+        for b in 0..(1 << (nv - 1)) {
+            res[b] = data[b << 1] * one_minus_p + data[(b << 1) + 1] * point;
+        }
+    }
+
+    res
+}
+
+pub fn evaluate_no_par<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {
+    assert_eq!(poly.num_vars, point.len());
+    fix_variables_no_par(poly, point).evaluations[0]
+}
+
+fn fix_variables_no_par<F: Field>(
+    poly: &DenseMultilinearExtension<F>,
+    partial_point: &[F],
+) -> DenseMultilinearExtension<F> {
+    assert!(
+        partial_point.len() <= poly.num_vars,
+        "invalid size of partial point"
+    );
+    let nv = poly.num_vars;
+    let mut poly = poly.evaluations.to_vec();
+    let dim = partial_point.len();
+    // evaluate single variable of partial point from left to right
+    for i in 1..dim + 1 {
+        let r = partial_point[i - 1];
+        let one_minus_r = F::one() - r;
+        for b in 0..(1 << (nv - i)) {
+            poly[b] = poly[b << 1] * one_minus_r + poly[(b << 1) + 1] * r;
+        }
+    }
+    DenseMultilinearExtension::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
+}
+
+/// merge a set of polynomials. Returns an error if the
+/// polynomials do not share a same number of nvs.
+pub fn merge_polynomials<F: PrimeField>(
+    polynomials: &[Rc<DenseMultilinearExtension<F>>],
+) -> Result<Rc<DenseMultilinearExtension<F>>, ArithErrors> {
+    let nv = polynomials[0].num_vars();
+    for poly in polynomials.iter() {
+        if nv != poly.num_vars() {
+            return Err(ArithErrors::InvalidParameters(
+                "num_vars do not match for polynomials".to_string(),
+            ));
+        }
+    }
+
+    let merged_nv = get_batched_nv(nv, polynomials.len());
+    let mut scalars = vec![];
+    for poly in polynomials.iter() {
+        scalars.extend_from_slice(poly.to_evaluations().as_slice());
+    }
+    scalars.extend_from_slice(vec![F::zero(); (1 << merged_nv) - scalars.len()].as_ref());
+    Ok(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        merged_nv, scalars,
+    )))
+}

arithmetic/src/univariate_polynomial.rs

@@ -0,0 +1,340 @@
+// TODO: remove
+#![allow(dead_code)]
+
+use crate::{bit_decompose, ArithErrors};
+use ark_ff::PrimeField;
+use ark_poly::{
+    univariate::DensePolynomial, EvaluationDomain, Evaluations, Radix2EvaluationDomain,
+};
+use ark_std::log2;
+
+/// Given a list of points, build `l(points)` which is a list of univariate
+/// polynomials that goes through the points; extend the dimension of the points
+/// by `log(points.len())` if `with_suffix` is set.
+pub fn build_l<F: PrimeField>(
+    points: &[Vec<F>],
+    domain: &Radix2EvaluationDomain<F>,
+    with_suffix: bool,
+) -> Result<Vec<DensePolynomial<F>>, ArithErrors> {
+    let mut uni_polys = Vec::new();
+    if with_suffix {
+        // 1.1 build the indexes and the univariate polys that go through the indexes
+        let prefix_len = log2(points.len()) as usize;
+        let indexes: Vec<Vec<bool>> = (0..points.len())
+            .map(|x| bit_decompose(x as u64, prefix_len))
+            .collect();
+        for i in 0..prefix_len {
+            let eval: Vec<F> = indexes
+                .iter()
+                .map(|x| F::from(x[prefix_len - i - 1]))
+                .collect();
+
+            uni_polys.push(Evaluations::from_vec_and_domain(eval, *domain).interpolate());
+        }
+    }
+    // 1.2 build the actual univariate polys that go through the points
+    uni_polys.extend_from_slice(build_l_internal(points, domain)?.as_slice());
+
+    Ok(uni_polys)
+}
+
+/// Given a list of points, build `l(points)` which is a list of univariate
+/// polynomials that goes through the points.
+pub(crate) fn build_l_internal<F: PrimeField>(
+    points: &[Vec<F>],
+    domain: &Radix2EvaluationDomain<F>,
+) -> Result<Vec<DensePolynomial<F>>, ArithErrors> {
+    let mut uni_polys = Vec::new();
+    let num_var = points[0].len();
+    // build the actual univariate polys that go through the points
+    for i in 0..num_var {
+        let mut eval: Vec<F> = points.iter().map(|x| x[i]).collect();
+        eval.extend_from_slice(vec![F::zero(); domain.size as usize - eval.len()].as_slice());
+        uni_polys.push(Evaluations::from_vec_and_domain(eval, *domain).interpolate())
+    }
+    Ok(uni_polys)
+}
+
+/// get the domain for the univariate polynomial
+#[inline]
+pub fn get_uni_domain<F: PrimeField>(
+    uni_poly_degree: usize,
+) -> Result<Radix2EvaluationDomain<F>, ArithErrors> {
+    let domain = match Radix2EvaluationDomain::<F>::new(uni_poly_degree) {
+        Some(p) => p,
+        None => {
+            return Err(ArithErrors::InvalidParameters(
+                "failed to build radix 2 domain".to_string(),
+            ))
+        },
+    };
+    Ok(domain)
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use ark_bls12_381::Fr;
+    use ark_ff::{field_new, One};
+    use ark_poly::UVPolynomial;
+
+    #[test]
+    fn test_build_l_with_suffix() -> Result<(), ArithErrors> {
+        test_build_l_with_suffix_helper::<Fr>()
+    }
+
+    fn test_build_l_with_suffix_helper<F: PrimeField>() -> Result<(), ArithErrors> {
+        // point 1 is [1, 2]
+        let point1 = vec![Fr::from(1u64), Fr::from(2u64)];
+
+        // point 2 is [3, 4]
+        let point2 = vec![Fr::from(3u64), Fr::from(4u64)];
+
+        // point 3 is [5, 6]
+        let point3 = vec![Fr::from(5u64), Fr::from(6u64)];
+
+        {
+            let domain = get_uni_domain::<Fr>(2)?;
+            let l = build_l(&[point1.clone(), point2.clone()], &domain, true)?;
+
+            // roots: [1, -1]
+            // l0 = -1/2 * x + 1/2
+            // l1 = -x + 2
+            // l2 = -x + 3
+            let l0 = DensePolynomial::from_coefficients_vec(vec![
+                Fr::one() / Fr::from(2u64),
+                -Fr::one() / Fr::from(2u64),
+            ]);
+            let l1 = DensePolynomial::from_coefficients_vec(vec![Fr::from(2u64), -Fr::one()]);
+            let l2 = DensePolynomial::from_coefficients_vec(vec![Fr::from(3u64), -Fr::one()]);
+
+            assert_eq!(l0, l[0], "l0 not equal");
+            assert_eq!(l1, l[1], "l1 not equal");
+            assert_eq!(l2, l[2], "l2 not equal");
+        }
+
+        {
+            let domain = get_uni_domain::<Fr>(3)?;
+            let l = build_l(&[point1, point2, point3], &domain, true)?;
+
+            // sage: q = 52435875175126190479447740508185965837690552500527637822603658699938581184513
+            // sage: P.<x> = PolynomialRing(Zmod(q))
+            // sage: root1 = 1
+            // sage: root2 = 0x8D51CCCE760304D0EC030002760300000001000000000000
+            // sage: root3 = -1
+            // sage: root4 = -root2
+            // Arkwork's code is a bit wired: it also interpolate (root4, 0)
+            // which returns a degree 3 polynomial, instead of degree 2
+
+            // ========================
+            // l0: [0, 0, 1]
+            // ========================
+            // sage: points = [(root1, 0), (root2, 0), (root3, 1), (root4, 0)]
+            // sage: P.lagrange_polynomial(points)
+            // 13108968793781547619861935127046491459422638125131909455650914674984645296128*x^3 +
+            // 39326906381344642859585805381139474378267914375395728366952744024953935888385*x^2 +
+            // 13108968793781547619861935127046491459422638125131909455650914674984645296128*x +
+            // 39326906381344642859585805381139474378267914375395728366952744024953935888385
+            let l0 = DensePolynomial::from_coefficients_vec(vec![
+                field_new!(
+                    Fr,
+                    "39326906381344642859585805381139474378267914375395728366952744024953935888385"
+                ),
+                field_new!(
+                    Fr,
+                    "13108968793781547619861935127046491459422638125131909455650914674984645296128"
+                ),
+                field_new!(
+                    Fr,
+                    "39326906381344642859585805381139474378267914375395728366952744024953935888385"
+                ),
+                field_new!(
+                    Fr,
+                    "13108968793781547619861935127046491459422638125131909455650914674984645296128"
+                ),
+            ]);
+
+            // ========================
+            // l1: [0, 1, 0]
+            // ========================
+            // sage: points = [(root1, 0), (root2, 1), (root3, 0), (root4, 0)]
+            // sage: P.lagrange_polynomial(points)
+            // 866286206518413079694067382671935694567563117191340490752*x^3 +
+            // 13108968793781547619861935127046491459422638125131909455650914674984645296128*x^2 +
+            // 52435875175126190478581454301667552757996485117855702128036095582747240693761*x +
+            // 39326906381344642859585805381139474378267914375395728366952744024953935888385
+            let l1 = DensePolynomial::from_coefficients_vec(vec![
+                field_new!(
+                    Fr,
+                    "39326906381344642859585805381139474378267914375395728366952744024953935888385"
+                ),
+                field_new!(
+                    Fr,
+                    "52435875175126190478581454301667552757996485117855702128036095582747240693761"
+                ),
+                field_new!(
+                    Fr,
+                    "13108968793781547619861935127046491459422638125131909455650914674984645296128"
+                ),
+                field_new!(
+                    Fr,
+                    "866286206518413079694067382671935694567563117191340490752"
+                ),
+            ]);
+
+            // ========================
+            // l2: [1, 3, 5]
+            // ========================
+            // sage: points = [(root1, 1), (root2, 3), (root3, 5), (root4, 0)]
+            // sage: P.lagrange_polynomial(points)
+            // 2598858619555239239082202148015807083702689351574021472255*x^3 +
+            // 13108968793781547619861935127046491459422638125131909455650914674984645296129*x^2 +
+            // 52435875175126190476848881888630726598608350352511830738900969348364559712256*x +
+            // 39326906381344642859585805381139474378267914375395728366952744024953935888387
+            let l2 = DensePolynomial::from_coefficients_vec(vec![
+                field_new!(
+                    Fr,
+                    "39326906381344642859585805381139474378267914375395728366952744024953935888387"
+                ),
+                field_new!(
+                    Fr,
+                    "52435875175126190476848881888630726598608350352511830738900969348364559712256"
+                ),
+                field_new!(
+                    Fr,
+                    "13108968793781547619861935127046491459422638125131909455650914674984645296129"
+                ),
+                field_new!(
+                    Fr,
+                    "2598858619555239239082202148015807083702689351574021472255"
+                ),
+            ]);
+
+            // ========================
+            // l3: [2, 4, 6]
+            // ========================
+            // sage: points = [(root1, 2), (root2, 4), (root3, 6), (root4, 0)]
+            // sage: P.lagrange_polynomial(points)
+            // 3465144826073652318776269530687742778270252468765361963007*x^3 +
+            // x^2 +
+            // 52435875175126190475982595682112313518914282969839895044333406231173219221504*x +
+            // 3
+            let l3 = DensePolynomial::from_coefficients_vec(vec![
+                Fr::from(3u64),
+                field_new!(
+                    Fr,
+                    "52435875175126190475982595682112313518914282969839895044333406231173219221504"
+                ),
+                Fr::one(),
+                field_new!(
+                    Fr,
+                    "3465144826073652318776269530687742778270252468765361963007"
+                ),
+            ]);
+
+            assert_eq!(l0, l[0], "l0 not equal");
+            assert_eq!(l1, l[1], "l1 not equal");
+            assert_eq!(l2, l[2], "l2 not equal");
+            assert_eq!(l3, l[3], "l3 not equal");
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_build_l() -> Result<(), ArithErrors> {
+        test_build_l_helper::<Fr>()
+    }
+
+    fn test_build_l_helper<F: PrimeField>() -> Result<(), ArithErrors> {
+        // point 1 is [1, 2]
+        let point1 = vec![Fr::from(1u64), Fr::from(2u64)];
+
+        // point 2 is [3, 4]
+        let point2 = vec![Fr::from(3u64), Fr::from(4u64)];
+
+        // point 3 is [5, 6]
+        let point3 = vec![Fr::from(5u64), Fr::from(6u64)];
+
+        {
+            let domain = get_uni_domain::<Fr>(2)?;
+            let l = build_l(&[point1.clone(), point2.clone()], &domain, false)?;
+
+            // roots: [1, -1]
+            // l0 = -x + 2
+            // l1 = -x + 3
+            let l0 = DensePolynomial::from_coefficients_vec(vec![Fr::from(2u64), -Fr::one()]);
+            let l1 = DensePolynomial::from_coefficients_vec(vec![Fr::from(3u64), -Fr::one()]);
+
+            assert_eq!(l0, l[0], "l0 not equal");
+            assert_eq!(l1, l[1], "l1 not equal");
+        }
+
+        {
+            let domain = get_uni_domain::<Fr>(3)?;
+            let l = build_l(&[point1, point2, point3], &domain, false)?;
+
+            // sage: q = 52435875175126190479447740508185965837690552500527637822603658699938581184513
+            // sage: P.<x> = PolynomialRing(Zmod(q))
+            // sage: root1 = 1
+            // sage: root2 = 0x8D51CCCE760304D0EC030002760300000001000000000000
+            // sage: root3 = -1
+            // sage: root4 = -root2
+            // Arkwork's code is a bit wired: it also interpolate (root4, 0)
+            // which returns a degree 3 polynomial, instead of degree 2
+
+            // ========================
+            // l0: [1, 3, 5]
+            // ========================
+            // sage: points = [(root1, 1), (root2, 3), (root3, 5), (root4, 0)]
+            // sage: P.lagrange_polynomial(points)
+            // 2598858619555239239082202148015807083702689351574021472255*x^3 +
+            // 13108968793781547619861935127046491459422638125131909455650914674984645296129*x^2 +
+            // 52435875175126190476848881888630726598608350352511830738900969348364559712256*x +
+            // 39326906381344642859585805381139474378267914375395728366952744024953935888387
+            let l0 = DensePolynomial::from_coefficients_vec(vec![
+                field_new!(
+                    Fr,
+                    "39326906381344642859585805381139474378267914375395728366952744024953935888387"
+                ),
+                field_new!(
+                    Fr,
+                    "52435875175126190476848881888630726598608350352511830738900969348364559712256"
+                ),
+                field_new!(
+                    Fr,
+                    "13108968793781547619861935127046491459422638125131909455650914674984645296129"
+                ),
+                field_new!(
+                    Fr,
+                    "2598858619555239239082202148015807083702689351574021472255"
+                ),
+            ]);
+
+            // ========================
+            // l1: [2, 4, 6]
+            // ========================
+            // sage: points = [(root1, 2), (root2, 4), (root3, 6), (root4, 0)]
+            // sage: P.lagrange_polynomial(points)
+            // 3465144826073652318776269530687742778270252468765361963007*x^3 +
+            // x^2 +
+            // 52435875175126190475982595682112313518914282969839895044333406231173219221504*x +
+            // 3
+            let l1 = DensePolynomial::from_coefficients_vec(vec![
+                Fr::from(3u64),
+                field_new!(
+                    Fr,
+                    "52435875175126190475982595682112313518914282969839895044333406231173219221504"
+                ),
+                Fr::one(),
+                field_new!(
+                    Fr,
+                    "3465144826073652318776269530687742778270252468765361963007"
+                ),
+            ]);
+
+            assert_eq!(l0, l[0], "l0 not equal");
+            assert_eq!(l1, l[1], "l1 not equal");
+        }
+        Ok(())
+    }
+}

arithmetic/src/util.rs

@@ -0,0 +1,96 @@
+use ark_ff::PrimeField;
+use ark_std::log2;
+
+/// Decompose an integer into a binary vector in little endian.
+pub 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
+}
+
+/// given the evaluation input `point` of the `index`-th polynomial,
+/// obtain the evaluation point in the merged polynomial
+pub fn gen_eval_point<F: PrimeField>(index: usize, index_len: usize, point: &[F]) -> Vec<F> {
+    let index_vec: Vec<F> = bit_decompose(index as u64, index_len)
+        .into_iter()
+        .map(|x| F::from(x))
+        .collect();
+    [point, &index_vec].concat()
+}
+
+/// Return the number of variables that one need for an MLE to
+/// batch the list of MLEs
+#[inline]
+pub fn get_batched_nv(num_var: usize, polynomials_len: usize) -> usize {
+    num_var + log2(polynomials_len) as usize
+}
+
+// Input index
+// - `i := (i_0, ...i_{n-1})`,
+// - `num_vars := n`
+// return three elements:
+// - `x0 := (i_1, ..., i_{n-1}, 0)`
+// - `x1 := (i_1, ..., i_{n-1}, 1)`
+// - `sign := i_0`
+#[inline]
+pub fn get_index(i: usize, num_vars: usize) -> (usize, usize, bool) {
+    let bit_sequence = bit_decompose(i as u64, num_vars);
+
+    // the last bit comes first here because of LE encoding
+    let x0 = project(&[[false].as_ref(), bit_sequence[..num_vars - 1].as_ref()].concat()) as usize;
+    let x1 = project(&[[true].as_ref(), bit_sequence[..num_vars - 1].as_ref()].concat()) as usize;
+
+    (x0, x1, bit_sequence[num_vars - 1])
+}
+
+/// Project a little endian binary vector into an integer.
+#[inline]
+pub(crate) fn project(input: &[bool]) -> u64 {
+    let mut res = 0;
+    for &e in input.iter().rev() {
+        res <<= 1;
+        res += e as u64;
+    }
+    res
+}
+
+#[cfg(test)]
+mod test {
+    use super::{bit_decompose, get_index, project};
+    use ark_std::{rand::RngCore, test_rng};
+
+    #[test]
+    fn test_decomposition() {
+        let mut rng = test_rng();
+        for _ in 0..100 {
+            let t = rng.next_u64();
+            let b = bit_decompose(t, 64);
+            let r = project(&b);
+            assert_eq!(t, r)
+        }
+    }
+
+    #[test]
+    fn test_get_index() {
+        let a = 0b1010;
+        let (x0, x1, sign) = get_index(a, 4);
+        assert_eq!(x0, 0b0100);
+        assert_eq!(x1, 0b0101);
+        assert!(sign);
+
+        let (x0, x1, sign) = get_index(a, 5);
+        assert_eq!(x0, 0b10100);
+        assert_eq!(x1, 0b10101);
+        assert!(!sign);
+
+        let a = 0b1111;
+        let (x0, x1, sign) = get_index(a, 4);
+        assert_eq!(x0, 0b1110);
+        assert_eq!(x1, 0b1111);
+        assert!(sign);
+    }
+}

arithmetic/src/virtual_polynomial.rs

@@ -1,7 +1,7 @@
 //! This module defines our main mathematical object `VirtualPolynomial`; and
 //! various functions associated with it.
 
-use crate::{errors::ArithErrors, multilinear_polynomial::random_zero_mle_list};
+use crate::{errors::ArithErrors, multilinear_polynomial::random_zero_mle_list, random_mle_list};
 use ark_ff::PrimeField;
 use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
 use ark_serialize::{CanonicalSerialize, SerializationError, Write};
@@ -324,42 +324,6 @@ impl<F: PrimeField> VirtualPolynomial<F> {
     }
 }
 
-/// Sample a random list of multilinear polynomials.
-/// Returns
-/// - the list of polynomials,
-/// - its sum of polynomial evaluations over the boolean hypercube.
-fn random_mle_list<F: PrimeField, R: RngCore>(
-    nv: usize,
-    degree: usize,
-    rng: &mut R,
-) -> (Vec<Rc<DenseMultilinearExtension<F>>>, F) {
-    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))
-    }
-    let mut sum = F::zero();
-
-    for _ in 0..(1 << nv) {
-        let mut product = F::one();
-
-        for e in multiplicands.iter_mut() {
-            let val = F::rand(rng);
-            e.push(val);
-            product *= val;
-        }
-        sum += product;
-    }
-
-    let list = multiplicands
-        .into_iter()
-        .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
-        .collect();
-
-    end_timer!(start);
-    (list, sum)
-}
-
 // This function build the eq(x, r) polynomial for any given r.
 //
 // Evaluate