diff --git a/spartan_parallel/src/custom_dense_mlpoly.rs b/spartan_parallel/src/custom_dense_mlpoly.rs index e04185e4..52ce8010 100644 --- a/spartan_parallel/src/custom_dense_mlpoly.rs +++ b/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 DensePolynomialPqx { // 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>>>, + mut z_mat: Vec>>>, num_proofs: Vec, max_num_proofs: usize, num_inputs: Vec, 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 DensePolynomialPqx { num_witness_secs: num_witness_secs.next_power_of_two(), num_inputs, max_num_inputs, - Z, + Z: z_mat, } } diff --git a/spartan_parallel/src/r1csinstance.rs b/spartan_parallel/src/r1csinstance.rs index bb92c4d0..73c88b75 100644 --- a/spartan_parallel/src/r1csinstance.rs +++ b/spartan_parallel/src/r1csinstance.rs @@ -285,21 +285,21 @@ impl R1CSInstance { ( 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(), diff --git a/spartan_parallel/src/r1csproof.rs b/spartan_parallel/src/r1csproof.rs index c834b69e..7cdc7fa0 100644 --- a/spartan_parallel/src/r1csproof.rs +++ b/spartan_parallel/src/r1csproof.rs @@ -311,7 +311,7 @@ impl R1CSProof { 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 R1CSProof { 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(),