From dcf68949b31b1d4a4446972e88ed317af1772d1b Mon Sep 17 00:00:00 2001
From: Dankrad Feist <mail@dankradfeist.de>
Date: Sat, 17 Sep 2022 18:00:54 +0100
Subject: [PATCH] Optimized eval_poly_l with batch inversion

---
 src/c_kzg.h   |  2 +-
 src/poly.c    | 29 +++++++++++++++++++++--------
 src/utility.h |  1 +
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/src/c_kzg.h b/src/c_kzg.h
index cf660e0..a3488aa 100644
--- a/src/c_kzg.h
+++ b/src/c_kzg.h
@@ -92,7 +92,7 @@ typedef struct {
 } poly_l; // Lagrange form
 
 void eval_poly(fr_t *out, const poly *p, const fr_t *x);
-void eval_poly_l(fr_t *out, const poly_l *p, const fr_t *x, const FFTSettings *fs);
+C_KZG_RET eval_poly_l(fr_t *out, const poly_l *p, const fr_t *x, const FFTSettings *fs);
 C_KZG_RET poly_inverse(poly *out, poly *b);
 C_KZG_RET poly_mul(poly *out, const poly *a, const poly *b);
 C_KZG_RET poly_mul_(poly *out, const poly *a, const poly *b, FFTSettings *fs);
diff --git a/src/poly.c b/src/poly.c
index 34346ed..d05c501 100644
--- a/src/poly.c
+++ b/src/poly.c
@@ -102,21 +102,32 @@ void eval_poly(fr_t *out, const poly *p, const fr_t *x) {
     }
 }
 
-// TODO: optimize via batch inversion
-void eval_poly_l(fr_t *out, const poly_l *p, const fr_t *x, const FFTSettings *fs) {
-  fr_t tmp;
+/**
+ * Evaluate a polynomial in Lagrange formover the finite field at a point.
+ *
+ * @param[out] out The value of the polynomial at the point @p x
+ * @param[in]  p   The polynomial in Lagrange form
+ * @param[in]  x   The x-coordinate to be evaluated
+ */
+C_KZG_RET eval_poly_l(fr_t *out, const poly_l *p, const fr_t *x, const FFTSettings *fs) {
+  fr_t tmp, *inverses_in, *inverses;
   uint64_t i;
   const uint64_t stride = fs->max_width / p->length;
 
-  *out = fr_zero;
+  TRY(new_fr_array(&inverses_in, p->length));
+  TRY(new_fr_array(&inverses, p->length));
   for (i = 0; i < p->length; i++) {
     if (fr_equal(x, &fs->expanded_roots_of_unity[i * stride])) {
       *out = p->values[i];
-      return;
+      return C_KZG_OK;
     }
-    fr_sub(&tmp, x, &fs->expanded_roots_of_unity[i * stride]);
-    fr_inv(&tmp, &tmp);
-    fr_mul(&tmp, &tmp, &fs->expanded_roots_of_unity[i * stride]);
+    fr_sub(&inverses_in[i], x, &fs->expanded_roots_of_unity[i * stride]);
+  }
+  TRY(fr_batch_inv(inverses, inverses_in, p->length));
+
+  *out = fr_zero;
+  for (i = 0; i < p->length; i++) {
+    fr_mul(&tmp, &inverses[i], &fs->expanded_roots_of_unity[i * stride]);
     fr_mul(&tmp, &tmp, &p->values[i]);
     fr_add(out, out, &tmp);
   }
@@ -125,6 +136,8 @@ void eval_poly_l(fr_t *out, const poly_l *p, const fr_t *x, const FFTSettings *f
   fr_pow(&tmp, x, p->length);
   fr_sub(&tmp, &tmp, &fr_one);
   fr_mul(out, out, &tmp);
+
+  return C_KZG_OK;
 }
 
 /**
diff --git a/src/utility.h b/src/utility.h
index e5af837..9803cb6 100644
--- a/src/utility.h
+++ b/src/utility.h
@@ -49,3 +49,4 @@ uint64_t next_power_of_two(uint64_t v);
 uint32_t reverse_bits(uint32_t a);
 uint32_t reverse_bits_limited(uint32_t n, uint32_t value);
 C_KZG_RET reverse_bit_order(void *values, size_t size, uint64_t n);
+C_KZG_RET fr_batch_inv(fr_t *out, const fr_t *a, size_t len);