From f3fffecd1ce7e8b6620cd5bac50c660efc20e48c Mon Sep 17 00:00:00 2001
From: Justin Traglia <95511699+jtraglia@users.noreply.github.com>
Date: Tue, 10 Oct 2023 09:07:44 -0500
Subject: [PATCH] Fix edge case in expand_root_of_unity (#375)

---
 src/c_kzg_4844.c | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/c_kzg_4844.c b/src/c_kzg_4844.c
index 2f65eb3..03c19b1 100644
--- a/src/c_kzg_4844.c
+++ b/src/c_kzg_4844.c
@@ -1557,13 +1557,16 @@ static C_KZG_RET bit_reversal_permutation(
 static C_KZG_RET expand_root_of_unity(
     fr_t *out, const fr_t *root, uint64_t width
 ) {
+    uint64_t i;
+    CHECK(width >= 2);
     out[0] = FR_ONE;
     out[1] = *root;
 
-    for (uint64_t i = 2; !fr_is_one(&out[i - 1]); i++) {
-        CHECK(i <= width);
+    for (i = 2; i <= width; i++) {
         blst_fr_mul(&out[i], &out[i - 1], root);
+        if (fr_is_one(&out[i])) break;
     }
+    CHECK(i == width);
     CHECK(fr_is_one(&out[width]));
 
     return C_KZG_OK;