Handle many missing (step 1)

src/zero_poly.c

@@ -112,11 +112,23 @@ C_KZG_RET pad_p(fr_t *out, uint64_t out_len, const fr_t *p, uint64_t p_len) {
 C_KZG_RET reduce_leaves(fr_t *dst, uint64_t len_dst, fr_t *scratch, uint64_t len_scratch, blst_fr **ps, uint64_t len_ps,
                         const uint64_t *len_p, const FFTSettings *fs) {
     CHECK(is_power_of_two(len_dst));
+    CHECK(len_scratch >= 3 * len_dst);
     CHECK(len_ps > 0);
     // The degree of the output is the sum of the degrees of the input polynomials.
     // TODO A more relaxed check should be ok: `len_ps * (len_p[0] - 1) < len_dst` (or even sum up the lengths)
-    CHECK(len_ps * len_p[0] <= len_dst);
+    // CHECK(len_ps * len_p[0] <= len_dst);
-    CHECK(len_scratch >= 3 * len_dst);
+    uint64_t total_length = 0;
+    for (int i = 0; i < len_ps; i++) {
+        total_length += len_p[i] - 1;
+    }
+    if (total_length + 1 > len_dst) {
+        printf("Total length: %lu, len dest: %lu\n", total_length, len_dst);
+        printf("\n");
+        for (int i = 0; i < len_ps; i++) {
+            printf("Len %d = %lu\n", i, len_p[i]);
+        }
+    }
+    CHECK(total_length + 1 <= len_dst);
 
     // Split `scratch` up into three equally sized working arrays
     fr_t *p_padded = scratch;
@@ -180,6 +192,7 @@ C_KZG_RET zero_polynomial_via_multiplication(fr_t *zero_eval, fr_t *zero_poly, u
         }
         return C_KZG_OK;
     }
+    CHECK(len_missing < length); // The output would be larger than length otherwise
     CHECK(length <= fs->max_width);
     CHECK(is_power_of_two(length));
 
@@ -188,16 +201,17 @@ C_KZG_RET zero_polynomial_via_multiplication(fr_t *zero_eval, fr_t *zero_poly, u
     uint64_t domain_stride = fs->max_width / length;
     uint64_t leaf_count = (len_missing + per_leaf - 1) / per_leaf;
     uint64_t n = next_power_of_two(leaf_count * per_leaf_poly);
+    if (n > length) n = length;
 
     if (len_missing <= per_leaf) {
         TRY(do_zero_poly_mul_leaf(zero_poly, length, missing_indices, len_missing, domain_stride, fs));
         TRY(fft_fr(zero_eval, zero_poly, false, length, fs));
         *zero_poly_len = len_missing + 1;
     } else {
-        CHECK(n <= length);
-
         // Work space for reducing the leaves - `zero_poly` is large enough due to the above check, so use that.
-        fr_t *work = zero_poly;
+        // fr_t *work = zero_poly;
+        fr_t *work;
+        TRY(new_fr_array(&work, next_power_of_two(leaf_count * per_leaf_poly)));
 
         // Build the leaves.
 
@@ -213,20 +227,23 @@ C_KZG_RET zero_polynomial_via_multiplication(fr_t *zero_eval, fr_t *zero_poly, u
             if (end > max) end = max;
             leaves[i] = &work[out_offset];
             leaf_lengths[i] = per_leaf_poly;
-            TRY(do_zero_poly_mul_leaf(leaves[i], leaf_lengths[i], &missing_indices[offset], end - offset, domain_stride,
+            TRY(do_zero_poly_mul_leaf(leaves[i], per_leaf_poly, &missing_indices[offset], end - offset, domain_stride,
                                       fs));
             offset += per_leaf;
             out_offset += per_leaf_poly;
         }
+        // Adjust the length of the last leaf
+        // leaf_lengths[leaf_count - 1] = 1 + len_missing % per_leaf;
+        leaf_lengths[leaf_count - 1] = 1 + len_missing - (leaf_count - 1) * per_leaf;
 
         // Now reduce all the leaves to a single poly
 
-        int reduction_factor = 4; // must be a power of 2
+        int reduction_factor = 4; // must be a power of 2 (why?)
         TRY(new_fr_array(&scratch, n * 3));
         while (leaf_count > 1) {
             uint64_t reduced_count = (leaf_count + reduction_factor - 1) / reduction_factor;
             // All the leaves are the same length, except possibly the last leaf, but that's ok.
-            uint64_t leaf_size = leaf_lengths[0];
+            uint64_t leaf_size = next_power_of_two(leaf_lengths[0]);
             for (uint64_t i = 0; i < reduced_count; i++) {
                 uint64_t start = i * reduction_factor;
                 uint64_t end = start + reduction_factor;
@@ -238,18 +255,24 @@ C_KZG_RET zero_polynomial_via_multiplication(fr_t *zero_eval, fr_t *zero_poly, u
                 }
                 reduced = work + start * leaf_size;
                 uint64_t reduced_len = out_end - start * leaf_size;
+                if (reduced_len > length) reduced_len = length;
                 if (end > leaf_count) {
                     end = leaf_count;
                 }
                 uint64_t leaves_slice_len = end - start;
-                if (end > start + 1) {
+                if (leaves_slice_len > 1) {
                     TRY(reduce_leaves(reduced, reduced_len, scratch, n * 3, &leaves[start], leaves_slice_len,
                                       &leaf_lengths[start], fs));
-                    leaf_lengths[i] = reduced_len;
+                    // leaf_lengths[i] = reduced_len;
-                } else {
+                    // } else {
-                    leaf_lengths[i] = leaf_lengths[start];
+                    //     leaf_lengths[i] = leaf_lengths[start];
                 }
                 leaves[i] = reduced;
+                uint64_t total_length = 0;
+                for (int j = start; j < end; j++) {
+                    total_length += leaf_lengths[j] - 1;
+                }
+                leaf_lengths[i] = total_length + 1;
             }
             leaf_count = reduced_count;
         }
@@ -260,6 +283,7 @@ C_KZG_RET zero_polynomial_via_multiplication(fr_t *zero_eval, fr_t *zero_poly, u
         }
         TRY(fft_fr(zero_eval, zero_poly, false, length, fs));
 
+        free(work);
         free(leaves);
         free(leaf_lengths);
         free(scratch);

src/zero_poly_test.c

@@ -251,7 +251,8 @@ void zero_poly_random(void) {
             FFTSettings fs;
             TEST_CHECK(C_KZG_OK == new_fft_settings(&fs, scale));
 
-            uint64_t missing[fs.max_width];
+            uint64_t *missing;
+            TEST_CHECK(C_KZG_OK == new_uint64_array(&missing, fs.max_width));
             int len_missing = 0;
 
             for (int i = 0; i < fs.max_width; i++) {
@@ -283,7 +284,7 @@ void zero_poly_random(void) {
                 ret = TEST_CHECK(fr_is_zero(&out));
                 TEST_MSG("Failed for missing[%d] = %lu", i, missing[i]);
             }
-            TEST_MSG("Failed for scale %d", scale);
+            TEST_MSG("Failed for scale = %d, len_missing = %d, zero_poly_len = %lu", scale, len_missing, zero_poly_len);
 
             fr_t *zero_eval_fft;
             TEST_CHECK(C_KZG_OK == new_fr_array(&zero_eval_fft, fs.max_width));
@@ -295,6 +296,7 @@ void zero_poly_random(void) {
                 TEST_CHECK(fr_is_zero(&zero_eval_fft[i]));
             }
 
+            free(missing);
             free(zero_poly);
             free(zero_eval);
             free(zero_eval_fft);
@@ -303,6 +305,105 @@ void zero_poly_random(void) {
     }
 }
 
+void zero_poly_all_but_one(void) {
+    FFTSettings fs;
+    TEST_CHECK(C_KZG_OK == new_fft_settings(&fs, 8));
+
+    uint64_t *missing;
+    TEST_CHECK(C_KZG_OK == new_uint64_array(&missing, fs.max_width));
+
+    // All but the first are missing
+    for (int i = 0; i < fs.max_width - 1; i++) {
+        missing[i] = i + 1;
+    }
+    int len_missing = fs.max_width - 1;
+
+    fr_t *zero_eval, *zero_poly;
+    uint64_t zero_poly_len;
+    TEST_CHECK(C_KZG_OK == new_fr_array(&zero_eval, fs.max_width));
+    TEST_CHECK(C_KZG_OK == new_fr_array(&zero_poly, fs.max_width));
+    TEST_CHECK(C_KZG_OK == zero_polynomial_via_multiplication(zero_eval, zero_poly, &zero_poly_len, fs.max_width,
+                                                              missing, len_missing, &fs));
+
+    poly p;
+    p.length = zero_poly_len;
+    p.coeffs = zero_poly;
+    int ret = 0;
+    for (int i = 0; i < len_missing; i++) {
+        fr_t out;
+        eval_poly(&out, &p, &fs.expanded_roots_of_unity[missing[i]]);
+        ret = TEST_CHECK(fr_is_zero(&out));
+        TEST_MSG("Failed for missing[%d] = %lu", i, missing[i]);
+    }
+
+    fr_t *zero_eval_fft;
+    TEST_CHECK(C_KZG_OK == new_fr_array(&zero_eval_fft, fs.max_width));
+    TEST_CHECK(C_KZG_OK == fft_fr(zero_eval_fft, zero_eval, true, fs.max_width, &fs));
+    for (uint64_t i = 0; i < zero_poly_len; i++) {
+        TEST_CHECK(fr_equal(&zero_poly[i], &zero_eval_fft[i]));
+    }
+    for (uint64_t i = zero_poly_len; i < fs.max_width; i++) {
+        TEST_CHECK(fr_is_zero(&zero_eval_fft[i]));
+    }
+
+    free(missing);
+    free(zero_poly);
+    free(zero_eval);
+    free(zero_eval_fft);
+    free_fft_settings(&fs);
+}
+
+void zero_poly_252(void) {
+    FFTSettings fs;
+    TEST_CHECK(C_KZG_OK == new_fft_settings(&fs, 8));
+
+    uint64_t *missing;
+    TEST_CHECK(C_KZG_OK == new_uint64_array(&missing, fs.max_width));
+
+    // 252 are missing
+    int len_missing = 252;
+    for (int i = 0; i < len_missing; i++) {
+        missing[i] = i;
+    }
+
+    fr_t *zero_eval, *zero_poly;
+    uint64_t zero_poly_len;
+    TEST_CHECK(C_KZG_OK == new_fr_array(&zero_eval, fs.max_width));
+    TEST_CHECK(C_KZG_OK == new_fr_array(&zero_poly, fs.max_width));
+    TEST_CHECK(C_KZG_OK == zero_polynomial_via_multiplication(zero_eval, zero_poly, &zero_poly_len, fs.max_width,
+                                                              missing, len_missing, &fs));
+
+    TEST_CHECK(zero_poly_len == 253);
+    TEST_MSG("ZeroPolyLen: expected %d, got %lu", len_missing + 1, zero_poly_len);
+
+    poly p;
+    p.length = zero_poly_len;
+    p.coeffs = zero_poly;
+    int ret = 0;
+    for (int i = 0; i < len_missing; i++) {
+        fr_t out;
+        eval_poly(&out, &p, &fs.expanded_roots_of_unity[missing[i]]);
+        ret = TEST_CHECK(fr_is_zero(&out));
+        TEST_MSG("Failed for missing[%d] = %lu", i, missing[i]);
+    }
+
+    fr_t *zero_eval_fft;
+    TEST_CHECK(C_KZG_OK == new_fr_array(&zero_eval_fft, fs.max_width));
+    TEST_CHECK(C_KZG_OK == fft_fr(zero_eval_fft, zero_eval, true, fs.max_width, &fs));
+    for (uint64_t i = 0; i < zero_poly_len; i++) {
+        TEST_CHECK(fr_equal(&zero_poly[i], &zero_eval_fft[i]));
+    }
+    for (uint64_t i = zero_poly_len; i < fs.max_width; i++) {
+        TEST_CHECK(fr_is_zero(&zero_eval_fft[i]));
+    }
+
+    free(missing);
+    free(zero_poly);
+    free(zero_eval);
+    free(zero_eval_fft);
+    free_fft_settings(&fs);
+}
+
 TEST_LIST = {
     {"ZERO_POLY_TEST", title},
     {"test_reduce_leaves", test_reduce_leaves},
@@ -310,5 +411,7 @@ TEST_LIST = {
     {"reduce_leaves_random", reduce_leaves_random},
     {"zero_poly_known", zero_poly_known},
     {"zero_poly_random", zero_poly_random},
+    {"zero_poly_all_but_one", zero_poly_all_but_one},
+    {"zero_poly_252", zero_poly_252},
     {NULL, NULL} /* zero record marks the end of the list */
 };