Refactor SNARKs

arithmetic/src/lib.rs

@@ -7,8 +7,8 @@ mod virtual_polynomial;
 pub use errors::ArithErrors;
 pub use multilinear_polynomial::{
     evaluate_no_par, evaluate_opt, fix_last_variables, fix_last_variables_no_par, fix_variables,
-    identity_permutation_mle, merge_polynomials, random_mle_list, random_permutation_mle,
-    random_zero_mle_list, DenseMultilinearExtension,
+    identity_permutation, identity_permutation_mles, merge_polynomials, random_mle_list,
+    random_permutation, random_permutation_mles, 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};

arithmetic/src/multilinear_polynomial.rs

@@ -72,31 +72,58 @@ pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
     list
 }
 
-/// An MLE that represent an identity permutation: `f(index) \mapto index`
-pub fn identity_permutation_mle<F: PrimeField>(
+pub fn identity_permutation<F: PrimeField>(num_vars: usize, num_chunks: usize) -> Vec<F> {
+    let len = (num_chunks as u64) * (1u64 << num_vars);
+    (0..len).map(F::from).collect()
+}
+
+/// A list of MLEs that represents an identity permutation
+pub fn identity_permutation_mles<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,
-    ))
+    num_chunks: usize,
+) -> Vec<Rc<DenseMultilinearExtension<F>>> {
+    let mut res = vec![];
+    for i in 0..num_chunks {
+        let shift = (i * (1 << num_vars)) as u64;
+        let s_id_vec = (shift..shift + (1u64 << num_vars)).map(F::from).collect();
+        res.push(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+            num_vars, s_id_vec,
+        )));
+    }
+    res
 }
 
-/// An MLE that represent a random permutation
-pub fn random_permutation_mle<F: PrimeField, R: RngCore>(
+pub fn random_permutation<F: PrimeField, R: RngCore>(
     num_vars: usize,
+    num_chunks: usize,
     rng: &mut R,
-) -> Rc<DenseMultilinearExtension<F>> {
-    let len = 1u64 << num_vars;
+) -> Vec<F> {
+    let len = (num_chunks as u64) * (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,
-    ))
+    s_perm_vec
+}
+
+/// A list of MLEs that represent a random permutation
+pub fn random_permutation_mles<F: PrimeField, R: RngCore>(
+    num_vars: usize,
+    num_chunks: usize,
+    rng: &mut R,
+) -> Vec<Rc<DenseMultilinearExtension<F>>> {
+    let s_perm_vec = random_permutation(num_vars, num_chunks, rng);
+    let mut res = vec![];
+    let n = 1 << num_vars;
+    for i in 0..num_chunks {
+        res.push(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+            num_vars,
+            s_perm_vec[i * n..i * n + n].to_vec(),
+        )));
+    }
+    res
 }
 
 pub fn evaluate_opt<F: Field>(poly: &DenseMultilinearExtension<F>, point: &[F]) -> F {

hyperplonk/src/mock.rs

@@ -1,6 +1,5 @@
-use arithmetic::identity_permutation_mle;
+use arithmetic::identity_permutation;
 use ark_ff::PrimeField;
-use ark_poly::MultilinearExtension;
 use ark_std::{log2, test_rng};
 
 use crate::{
@@ -93,7 +92,7 @@ impl<F: PrimeField> MockCircuit<F> {
             gate_func: gate.clone(),
         };
 
-        let permutation = identity_permutation_mle(merged_nv as usize).to_evaluations();
+        let permutation = identity_permutation(merged_nv as usize, 1);
         let index = HyperPlonkIndex {
             params,
             permutation,