From e54a4a2790578160cd834bcf601b6e83b1a45dfe Mon Sep 17 00:00:00 2001
From: Danny Willems <be.danny.willems@gmail.com>
Date: Wed, 4 Sep 2024 15:59:13 -0700
Subject: [PATCH] Kimchi/LagrangeBasisEvaluations: optimize evaluate_boolean

---
 kimchi/src/lagrange_basis_evaluations.rs | 25 ++++++++++++++++++------
 kimchi/src/prover.rs                     |  4 ++--
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/kimchi/src/lagrange_basis_evaluations.rs b/kimchi/src/lagrange_basis_evaluations.rs
index 8ab338d66a..2b32f1be1e 100644
--- a/kimchi/src/lagrange_basis_evaluations.rs
+++ b/kimchi/src/lagrange_basis_evaluations.rs
@@ -102,21 +102,34 @@ impl<F: FftField> LagrangeBasisEvaluations<F> {
     /// polynomial at a point, assuming that the given evaluations are either `0`
     /// or `1` at every point of the domain.
     ///
+    /// A slight optimisation is performed here, where we only use addition on
+    /// integers (usize) instead of the addition on field elements. We convert
+    /// only at the end into a field element.
+    ///
     /// This method can particularly be useful when the polynomials represent
     /// (boolean) selectors in a circuit.
-    pub fn evaluate_boolean<D: EvaluationDomain<F>>(&self, p: &Evaluations<F, D>) -> Vec<F> {
+    ///
+    /// # Safety
+    ///
+    /// There is no check on the input, i.e. we do not verify the assumptions
+    /// that all evaluation points are `1` or `0`.
+    ///
+    /// There is also an additional assumption that the number of evaluation
+    /// points is not greater than 2^64 - 1 (i.e. the machine word size on
+    /// 64bits machine).
+    pub unsafe fn evaluate_boolean<D: EvaluationDomain<F>>(&self, p: &Evaluations<F, D>) -> Vec<F> {
         assert_eq!(p.evals.len() % self.domain_size(), 0);
         let stride = p.evals.len() / self.domain_size();
         self.evals
             .iter()
             .map(|evals| {
-                let mut result = F::zero();
-                for (i, e) in evals.iter().enumerate() {
+                let mut result: usize = 0;
+                for (i, _e) in evals.iter().enumerate() {
                     if !p.evals[stride * i].is_zero() {
-                        result += e;
+                        result += 1;
                     }
                 }
-                result
+                F::from(result as u64)
             })
             .collect()
     }
@@ -341,7 +354,7 @@ mod tests {
 
         let evaluator = LagrangeBasisEvaluations::new(domain.size(), domain, x);
 
-        let y = evaluator.evaluate_boolean(&evals);
+        let y = unsafe { evaluator.evaluate_boolean(&evals) };
         let expected = vec![evals.interpolate().evaluate(&x)];
         assert_eq!(y, expected)
     }
diff --git a/kimchi/src/prover.rs b/kimchi/src/prover.rs
index 547c6a2636..e54a153746 100644
--- a/kimchi/src/prover.rs
+++ b/kimchi/src/prover.rs
@@ -962,8 +962,8 @@ where
 
         let chunked_evals_for_selector =
             |p: &Evaluations<G::ScalarField, D<G::ScalarField>>| PointEvaluations {
-                zeta: zeta_evals.evaluate_boolean(p),
-                zeta_omega: zeta_omega_evals.evaluate_boolean(p),
+                zeta: unsafe { zeta_evals.evaluate_boolean(p) },
+                zeta_omega: unsafe { zeta_omega_evals.evaluate_boolean(p) },
             };
 
         let chunked_evals_for_evaluations =