From c9f68457155f771303409aced5601ba302f385ae Mon Sep 17 00:00:00 2001
From: Justin Traglia <jtraglia@pm.me>
Date: Wed, 18 Sep 2024 14:11:18 -0500
Subject: [PATCH] Clean up toeplitz_coeffs_stride

---
 src/eip7594/cell.h |  5 ++++-
 src/eip7594/fft.c  |  2 +-
 src/eip7594/fk20.c | 39 ++++++++++++++++-----------------------
 3 files changed, 21 insertions(+), 25 deletions(-)

diff --git a/src/eip7594/cell.h b/src/eip7594/cell.h
index f36ffb6c..57f4f734 100644
--- a/src/eip7594/cell.h
+++ b/src/eip7594/cell.h
@@ -30,12 +30,15 @@
 /** The number of bytes in a single cell. */
 #define BYTES_PER_CELL (FIELD_ELEMENTS_PER_CELL * BYTES_PER_FIELD_ELEMENT)
 
-/** The Reed-Solomon erasuring coding expansion factor. */
+/** The Reed-Solomon erasure coding expansion factor. */
 #define EXPANSION_FACTOR 2
 
 /** The number of field elements in an extended blob */
 #define FIELD_ELEMENTS_PER_EXT_BLOB (FIELD_ELEMENTS_PER_BLOB * EXPANSION_FACTOR)
 
+/** The number of cells in a blob. */
+#define CELLS_PER_BLOB (FIELD_ELEMENTS_PER_BLOB / FIELD_ELEMENTS_PER_CELL)
+
 /** The number of cells in an extended blob. */
 #define CELLS_PER_EXT_BLOB (FIELD_ELEMENTS_PER_EXT_BLOB / FIELD_ELEMENTS_PER_CELL)
 
diff --git a/src/eip7594/fft.c b/src/eip7594/fft.c
index 182f6385..49945b12 100644
--- a/src/eip7594/fft.c
+++ b/src/eip7594/fft.c
@@ -17,7 +17,7 @@
 #include "eip7594/fft.h"
 #include "common/alloc.h"
 #include "common/utils.h"
-#include "eip4844/blob.h"
+#include "eip7594/cell.h"
 #include "eip7594/poly.h"
 
 #include <string.h> /* For memcpy */
diff --git a/src/eip7594/fk20.c b/src/eip7594/fk20.c
index 6f304c48..a8027500 100644
--- a/src/eip7594/fk20.c
+++ b/src/eip7594/fk20.c
@@ -25,31 +25,27 @@
 /**
  * Reorder and extend polynomial coefficients for the toeplitz method, strided version.
  *
- * @param[out]  out     The reordered polynomial, size `n * 2 / stride`
- * @param[in]   in      The input polynomial, size `n`
- * @param[in]   n       The size of the input polynomial
+ * @param[out]  out     The reordered polynomial, length `CELLS_PER_EXT_BLOB`
+ * @param[in]   in      The input polynomial, length `FIELD_ELEMENTS_PER_BLOB`
  * @param[in]   offset  The offset
- * @param[in]   stride  The stride
  */
-static C_KZG_RET toeplitz_coeffs_stride(
-    fr_t *out, const fr_t *in, size_t n, size_t offset, size_t stride
-) {
-    size_t k, k2;
-
-    if (stride == 0) return C_KZG_BADARGS;
+static void toeplitz_coeffs_stride(fr_t *out, const fr_t *in, size_t offset) {
+    /* Calculate starting indices */
+    size_t out_start = CELLS_PER_BLOB + 2;
+    size_t in_start = CELLS_PER_EXT_BLOB - offset - 1;
 
-    k = n / stride;
-    k2 = k * 2;
+    /* Set the first element */
+    out[0] = in[FIELD_ELEMENTS_PER_BLOB - 1 - offset];
 
-    out[0] = in[n - 1 - offset];
-    for (size_t i = 1; i <= k + 1 && i < k2; i++) {
+    /* Initialize these elements to zero */
+    for (size_t i = 1; i < out_start; i++) {
         out[i] = FR_ZERO;
     }
-    for (size_t i = k + 2, j = 2 * stride - offset - 1; i < k2; i++, j += stride) {
-        out[i] = in[j];
-    }
 
-    return C_KZG_OK;
+    /* Copy elements with a fixed stride */
+    for (size_t i = 0; i < CELLS_PER_EXT_BLOB - out_start; i++) {
+        out[out_start + i] = in[in_start + i * FIELD_ELEMENTS_PER_CELL];
+    }
 }
 
 /**
@@ -77,7 +73,7 @@ C_KZG_RET compute_fk20_cell_proofs(g1_t *out, const fr_t *p, const KZGSettings *
 
     /* Initialize length variables */
     k = FIELD_ELEMENTS_PER_BLOB / FIELD_ELEMENTS_PER_CELL;
-    k2 = k * 2;
+    k2 = k * EXPANSION_FACTOR;
 
     /* Do allocations */
     ret = new_fr_array(&toeplitz_coeffs, k2);
@@ -112,10 +108,7 @@ C_KZG_RET compute_fk20_cell_proofs(g1_t *out, const fr_t *p, const KZGSettings *
 
     /* Compute toeplitz coefficients and organize by column */
     for (size_t i = 0; i < FIELD_ELEMENTS_PER_CELL; i++) {
-        ret = toeplitz_coeffs_stride(
-            toeplitz_coeffs, p, FIELD_ELEMENTS_PER_BLOB, i, FIELD_ELEMENTS_PER_CELL
-        );
-        if (ret != C_KZG_OK) goto out;
+        toeplitz_coeffs_stride(toeplitz_coeffs, p, i);
         ret = fr_fft(toeplitz_coeffs_fft, toeplitz_coeffs, k2, s);
         if (ret != C_KZG_OK) goto out;
         for (size_t j = 0; j < k2; j++) {