From ea18b23bf563f34c2ae9a6df0a15319e7ef0763f Mon Sep 17 00:00:00 2001
From: Ben Edgington <ben@benjaminion.xyz>
Date: Wed, 10 Feb 2021 10:39:50 +0000
Subject: [PATCH] Fix bug with strided FFT

---
 src/fft_fr.c      |  4 ++--
 src/fft_fr_test.c | 21 +++++++++++++++++++++
 src/fft_g1.c      |  4 ++--
 src/fft_g1_test.c | 18 ++++++++++++++++++
 4 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/fft_fr.c b/src/fft_fr.c
index 7d35366..ea9d561 100644
--- a/src/fft_fr.c
+++ b/src/fft_fr.c
@@ -60,12 +60,12 @@ C_KZG_RET fft_fr(blst_fr *out, const blst_fr *in, const FFTSettings *fs, bool in
         blst_fr inv_len;
         fr_from_uint64(&inv_len, n);
         blst_fr_eucl_inverse(&inv_len, &inv_len);
-        fft_fr_fast(out, in, 1, fs->reverse_roots_of_unity, stride, fs->max_width);
+        fft_fr_fast(out, in, 1, fs->reverse_roots_of_unity, stride, n);
         for (uint64_t i = 0; i < fs->max_width; i++) {
             blst_fr_mul(&out[i], &out[i], &inv_len);
         }
     } else {
-        fft_fr_fast(out, in, 1, fs->expanded_roots_of_unity, stride, fs->max_width);
+        fft_fr_fast(out, in, 1, fs->expanded_roots_of_unity, stride, n);
     }
     return C_KZG_OK;
 }
diff --git a/src/fft_fr_test.c b/src/fft_fr_test.c
index 09066bb..e66e5f0 100644
--- a/src/fft_fr_test.c
+++ b/src/fft_fr_test.c
@@ -109,10 +109,31 @@ void inverse_fft(void) {
     free_fft_settings(&fs);
 }
 
+void stride_fft(void) {
+    unsigned int size1 = 9, size2 = 12;
+    uint64_t width = size1 < size2 ? (uint64_t)1 << size1 : (uint64_t)1 << size2;
+    FFTSettings fs1, fs2;
+    TEST_CHECK(new_fft_settings(&fs1, size1) == C_KZG_OK);
+    TEST_CHECK(new_fft_settings(&fs2, size2) == C_KZG_OK);
+    blst_fr data[width], coeffs1[width], coeffs2[width];
+    for (int i = 0; i < width; i++) {
+        fr_from_uint64(data + i, i);
+    }
+
+    TEST_CHECK(fft_fr(coeffs1, data, &fs1, false, width) == C_KZG_OK);
+    TEST_CHECK(fft_fr(coeffs2, data, &fs2, false, width) == C_KZG_OK);
+
+    for (int i = 0; i < width; i++) {
+        TEST_CHECK(fr_equal(coeffs1 + i, coeffs2 + i));
+    }
+}
+
 TEST_LIST = {
     {"FFT_FR_TEST", title},
     {"compare_sft_fft", compare_sft_fft},
     {"roundtrip_fft", roundtrip_fft},
     {"inverse_fft", inverse_fft},
+    {"stride_fft", stride_fft},
+
     {NULL, NULL} /* zero record marks the end of the list */
 };
diff --git a/src/fft_g1.c b/src/fft_g1.c
index 4ccafda..bfe01f0 100644
--- a/src/fft_g1.c
+++ b/src/fft_g1.c
@@ -59,12 +59,12 @@ C_KZG_RET fft_g1(blst_p1 *out, blst_p1 *in, FFTSettings *fs, bool inv, uint64_t
         blst_fr inv_len;
         fr_from_uint64(&inv_len, n);
         blst_fr_eucl_inverse(&inv_len, &inv_len);
-        fft_g1_fast(out, in, 1, fs->reverse_roots_of_unity, stride, fs->max_width);
+        fft_g1_fast(out, in, 1, fs->reverse_roots_of_unity, stride, n);
         for (uint64_t i = 0; i < fs->max_width; i++) {
             p1_mul(&out[i], &out[i], &inv_len);
         }
     } else {
-        fft_g1_fast(out, in, 1, fs->expanded_roots_of_unity, stride, fs->max_width);
+        fft_g1_fast(out, in, 1, fs->expanded_roots_of_unity, stride, n);
     }
     return C_KZG_OK;
 }
diff --git a/src/fft_g1_test.c b/src/fft_g1_test.c
index dc0c7f9..00c9cdb 100644
--- a/src/fft_g1_test.c
+++ b/src/fft_g1_test.c
@@ -70,9 +70,27 @@ void roundtrip_fft(void) {
     free_fft_settings(&fs);
 }
 
+void stride_fft(void) {
+    unsigned int size1 = 9, size2 = 12;
+    uint64_t width = size1 < size2 ? (uint64_t)1 << size1 : (uint64_t)1 << size2;
+    FFTSettings fs1, fs2;
+    TEST_CHECK(new_fft_settings(&fs1, size1) == C_KZG_OK);
+    TEST_CHECK(new_fft_settings(&fs2, size2) == C_KZG_OK);
+    blst_p1 data[width], coeffs1[width], coeffs2[width];
+    make_data(data, width);
+
+    TEST_CHECK(fft_g1(coeffs1, data, &fs1, false, width) == C_KZG_OK);
+    TEST_CHECK(fft_g1(coeffs2, data, &fs2, false, width) == C_KZG_OK);
+
+    for (int i = 0; i < width; i++) {
+        TEST_CHECK(blst_p1_is_equal(coeffs1 + i, coeffs2 + i));
+    }
+}
+
 TEST_LIST = {
     {"FFT_G1_TEST", title},
     {"compare_sft_fft", compare_sft_fft},
     {"roundtrip_fft", roundtrip_fft},
+    {"stride_fft", stride_fft},
     {NULL, NULL} /* zero record marks the end of the list */
 };