From e9fec1591e5f1c909b8782403b0db86a20de91d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Garillot?= <francois@garillot.net>
Date: Thu, 8 Feb 2024 14:00:47 -0500
Subject: [PATCH] refactor: Enhance efficiency of generating power polynomials

- Introduced `#[must_use]` annotation to the `squares` function in `power.rs` to ensure the returned result is handled appropriately.
- avoid re-computing powers of tau over and over in batched_ppsnark
---
 src/spartan/batched_ppsnark.rs | 11 +++++++----
 src/spartan/polys/eq.rs        |  3 ++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/spartan/batched_ppsnark.rs b/src/spartan/batched_ppsnark.rs
index aabf6dc61..f6b2c42f1 100644
--- a/src/spartan/batched_ppsnark.rs
+++ b/src/spartan/batched_ppsnark.rs
@@ -192,6 +192,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
     // N[i] = max(|Aᵢ|+|Bᵢ|+|Cᵢ|, 2*num_varsᵢ, num_consᵢ)
     let N = pk.S_repr.iter().map(|s| s.N).collect::<Vec<_>>();
     assert!(N.iter().all(|&Ni| Ni.is_power_of_two()));
+    let N_max = *N.iter().max().unwrap();
 
     let num_instances = U.len();
 
@@ -247,13 +248,15 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
     // Compute eq(tau) for each instance in log2(Ni) variables
     let tau = transcript.squeeze(b"t")?;
+    let all_taus = PowPolynomial::squares(&tau, N_max.log_2());
+
     let (polys_tau, coords_tau): (Vec<_>, Vec<_>) = N
-      .iter()
+      .par_iter()
       .map(|&N_i| {
         let log_Ni = N_i.log_2();
-        let poly = PowPolynomial::new(&tau, log_Ni);
-        let evals = poly.evals();
-        let coords = poly.coordinates();
+        let eqp: EqPolynomial<_> = all_taus[..log_Ni].iter().cloned().collect();
+        let evals = eqp.evals();
+        let coords = eqp.r;
         (evals, coords)
       })
       .unzip();
diff --git a/src/spartan/polys/eq.rs b/src/spartan/polys/eq.rs
index 5933950f3..0db4edb4d 100644
--- a/src/spartan/polys/eq.rs
+++ b/src/spartan/polys/eq.rs
@@ -16,7 +16,7 @@ use rayon::prelude::{IndexedParallelIterator, IntoParallelRefMutIterator, Parall
 /// For instance, for e = 6 (with a binary representation of 0b110), the vector r would be [1, 1, 0].
 #[derive(Debug)]
 pub struct EqPolynomial<Scalar> {
-  pub(crate) r: Vec<Scalar>,
+  pub(in crate::spartan) r: Vec<Scalar>,
 }
 
 impl<Scalar: PrimeField> EqPolynomial<Scalar> {
@@ -43,6 +43,7 @@ impl<Scalar: PrimeField> EqPolynomial<Scalar> {
   /// Evaluates the `EqPolynomial` at all the `2^|r|` points in its domain.
   ///
   /// Returns a vector of Scalars, each corresponding to the polynomial evaluation at a specific point.
+  #[must_use = "this returns an expensive vector and leaves self unchanged"]
   pub fn evals(&self) -> Vec<Scalar> {
     Self::evals_from_points(&self.r)
   }