From 903a13fb98392e4df592854ae023cc3a8e2c84d7 Mon Sep 17 00:00:00 2001
From: George Kadianakis <desnacked@riseup.net>
Date: Thu, 16 Mar 2023 19:52:59 +0200
Subject: [PATCH] Simplify `fr_batch_inv` and reject zero inputs (#215)

---
 src/c_kzg_4844.c      | 55 +++++++++++++++++++++++++++----------------
 src/test_c_kzg_4844.c |  8 ++-----
 2 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/src/c_kzg_4844.c b/src/c_kzg_4844.c
index d0f146d..fc09504 100644
--- a/src/c_kzg_4844.c
+++ b/src/c_kzg_4844.c
@@ -254,6 +254,20 @@ static bool fr_is_one(const fr_t *p) {
     return a[0] == 1 && a[1] == 0 && a[2] == 0 && a[3] == 0;
 }
 
+/**
+ * Test whether the operand is zero in the finite field.
+ *
+ * @param[in] p The field element to be checked
+ *
+ * @retval true  The element is zero
+ * @retval false The element is not zero
+ */
+static bool fr_is_zero(const fr_t *p) {
+    uint64_t a[4];
+    blst_uint64_from_fr(a, p);
+    return a[0] == 0 && a[1] == 0 && a[2] == 0 && a[3] == 0;
+}
+
 /**
  * Test whether two field elements are equal.
  *
@@ -332,42 +346,43 @@ static void fr_from_uint64(fr_t *out, uint64_t n) {
 /**
  * Montgomery batch inversion in finite field.
  *
+ * @remark Return C_KZG_BADARGS if a zero is found in the input.
+ *
  * @remark This function does not support in-place computation (i.e. `a` MUST
  * NOT point to the same place as `out`)
  *
+ * @remark This function only supports len > 0.
+ *
  * @param[out] out The inverses of @p a, length @p len
  * @param[in]  a   A vector of field elements, length @p len
  * @param[in]  len The number of field elements
  */
-static C_KZG_RET fr_batch_inv(fr_t *out, const fr_t *a, size_t len) {
-    C_KZG_RET ret;
-    fr_t *prod = NULL;
-    fr_t inv;
-    size_t i;
+static C_KZG_RET fr_batch_inv(fr_t *out, const fr_t *a, int len) {
+    int i;
 
     assert(len > 0);
     assert(a != out);
 
-    ret = new_fr_array(&prod, len);
-    if (ret != C_KZG_OK) goto out;
+    fr_t accumulator = FR_ONE;
 
-    prod[0] = a[0];
-
-    for (i = 1; i < len; i++) {
-        blst_fr_mul(&prod[i], &a[i], &prod[i - 1]);
+    for (i = 0; i < len; i++) {
+        out[i] = accumulator;
+        blst_fr_mul(&accumulator, &accumulator, &a[i]);
     }
 
-    blst_fr_eucl_inverse(&inv, &prod[len - 1]);
-
-    for (i = len - 1; i > 0; i--) {
-        blst_fr_mul(&out[i], &inv, &prod[i - 1]);
-        blst_fr_mul(&inv, &a[i], &inv);
+    /* Bail on any zero input */
+    if (fr_is_zero(&accumulator)) {
+        return C_KZG_BADARGS;
     }
-    out[0] = inv;
 
-out:
-    c_kzg_free(prod);
-    return ret;
+    blst_fr_eucl_inverse(&accumulator, &accumulator);
+
+    for (i = len - 1; i >= 0; i--) {
+        blst_fr_mul(&out[i], &out[i], &accumulator);
+        blst_fr_mul(&accumulator, &accumulator, &a[i]);
+    }
+
+    return C_KZG_OK;
 }
 
 /**
diff --git a/src/test_c_kzg_4844.c b/src/test_c_kzg_4844.c
index 52b382f..8bd5dbf 100644
--- a/src/test_c_kzg_4844.c
+++ b/src/test_c_kzg_4844.c
@@ -287,6 +287,7 @@ static void test_fr_batch_inv__test_consistent(void) {
     }
 }
 
+/** Make sure that batch inverse doesn't support zeroes */
 static void test_fr_batch_inv__test_zero(void) {
     C_KZG_RET ret;
     fr_t a[32], batch_inverses[32];
@@ -298,12 +299,7 @@ static void test_fr_batch_inv__test_zero(void) {
     a[5] = FR_ZERO;
 
     ret = fr_batch_inv(batch_inverses, a, 32);
-    ASSERT_EQUALS(ret, C_KZG_OK);
-
-    for (size_t i = 0; i < 32; i++) {
-        bool ok = fr_equal(&batch_inverses[i], &FR_ZERO);
-        ASSERT_EQUALS(ok, true);
-    }
+    ASSERT_EQUALS(ret, C_KZG_BADARGS);
 }
 
 ///////////////////////////////////////////////////////////////////////////////