More efficient custom_dense_poly construction

spartan_parallel/src/custom_dense_mlpoly.rs

@@ -1,5 +1,6 @@
 #![allow(clippy::too_many_arguments)]
-use std::cmp::min;
+use std::cmp::{max, min};
+use std::mem;
 
 use crate::dense_mlpoly::DensePolynomial;
 use crate::scalar::SpartanExtensionField;
@@ -63,40 +64,51 @@ impl<S: SpartanExtensionField> DensePolynomialPqx<S> {
   // Assume z_mat is in its standard form of (p, q, x)
   // Reverse q and x and convert it to (p, q_rev, x_rev)
   pub fn new_rev(
-    z_mat: &Vec<Vec<Vec<Vec<S>>>>,
+    mut z_mat: Vec<Vec<Vec<Vec<S>>>>,
     num_proofs: Vec<usize>,
     max_num_proofs: usize,
     num_inputs: Vec<usize>,
     max_num_inputs: usize,
   ) -> Self {
-    let mut Z = Vec::new();
     let num_instances = z_mat.len();
     let num_witness_secs = z_mat[0][0].len();
     for p in 0..num_instances {
-      Z.push(vec![
-        vec![
-          vec![S::field_zero(); num_inputs[p]];
-          num_witness_secs
-        ];
-        num_proofs[p]
-      ]);
-
-      let step_q = max_num_proofs / num_proofs[p];
+      // Reverse the bits of x in place
       let step_x = max_num_inputs / num_inputs[p];
-      for q in 0..num_proofs[p] {
-        // Reverse the bits of q. q_rev is a multiple of step_q
-        let q_rev = rev_bits(q, max_num_proofs);
-        // Now q_rev is between 0 to num_proofs[p]
-        let q_rev = q_rev / step_q;
-
-        for x in 0..num_inputs[p] {
+      let mut x_swapped = vec![false; num_inputs[p]];
+      for x in 0..num_inputs[p] {
+        if !x_swapped[x] {
           // Reverse the bits of x. x_rev is a multiple of step_x
           let x_rev = rev_bits(x, max_num_inputs);
           // Now x_rev is between 0 to num_inputs[p]
           let x_rev = x_rev / step_x;
-          for w in 0..num_witness_secs {
-            Z[p][q_rev][w][x_rev] = z_mat[p][q][w][x];
+          for q in 0..num_proofs[p] {
+            for w in 0..num_witness_secs {
+              let tmp = z_mat[p][q][w][x];
+              z_mat[p][q][w][x] = z_mat[p][q][w][x_rev];
+              z_mat[p][q][w][x_rev] = tmp;
+            }
           }
+          x_swapped[x_rev] = true;
+        }
+      }
+      // Reverse the bits of q
+      let step_q = max_num_proofs / num_proofs[p];
+      let mut q_swapped = vec![false; num_proofs[p]];
+      for q in 0..num_proofs[p] {
+        if !q_swapped[q] {
+          // Reverse the bits of q. q_rev is a multiple of step_q
+          let q_rev = rev_bits(q, max_num_proofs);
+          // Now q_rev is between 0 to num_proofs[p]
+          let q_rev = q_rev / step_q;
+          if q != q_rev {
+            let q_low = min(q, q_rev);
+            let q_high = max(q, q_rev);
+            let (left, right) = z_mat[p].split_at_mut(q_low + 1);
+            mem::swap(&mut left[q_low], &mut right[q_high - q_low - 1]);
+          }
+
+          q_swapped[q_rev] = true;
         }
       }
     }
@@ -107,7 +119,7 @@ impl<S: SpartanExtensionField> DensePolynomialPqx<S> {
       num_witness_secs: num_witness_secs.next_power_of_two(),
       num_inputs,
       max_num_inputs,
-      Z,
+      Z: z_mat,
     }
   }
 

spartan_parallel/src/r1csinstance.rs

@@ -285,21 +285,21 @@ impl<S: SpartanExtensionField + Send + Sync> R1CSInstance<S> {
 
     (
       DensePolynomialPqx::new_rev(
-        &Az,
+        Az,
         num_proofs.clone(),
         max_num_proofs,
         num_cons.clone(),
         max_num_cons,
       ),
       DensePolynomialPqx::new_rev(
-        &Bz,
+        Bz,
         num_proofs.clone(),
         max_num_proofs,
         num_cons.clone(),
         max_num_cons,
       ),
       DensePolynomialPqx::new_rev(
-        &Cz,
+        Cz,
         num_proofs,
         max_num_proofs,
         num_cons.clone(),

spartan_parallel/src/r1csproof.rs

@@ -311,7 +311,7 @@ impl<S: SpartanExtensionField + Send + Sync> R1CSProof<S> {
       evals_ABC
     };
     let mut ABC_poly = DensePolynomialPqx::new_rev(
-      &evals_ABC,
+      evals_ABC,
       vec![1; num_instances],
       1,
       num_inputs.clone(),
@@ -322,7 +322,7 @@ impl<S: SpartanExtensionField + Send + Sync> R1CSProof<S> {
     let timer_tmp = Timer::new("prove_z_gen");
     // Construct a p * q * len(z) matrix Z and bound it to r_q
     let mut Z_poly = DensePolynomialPqx::new_rev(
-      &z_mat,
+      z_mat,
       num_proofs.clone(),
       max_num_proofs,
       num_inputs.clone(),