Improve coset factor computation (#7)

src/eip7594/eip7594.c

@@ -46,6 +46,9 @@ static const char *RANDOM_CHALLENGE_DOMAIN_VERIFY_CELL_KZG_PROOF_BATCH = "RCKZGC
  *
  * for (size_t i = 0; i < CELLS_PER_EXT_BLOB; i++)
  *   printf("%#04llx,\n", reverse_bits_limited(CELLS_PER_EXT_BLOB, i));
+ *
+ * Because of the way our evaluation domain is defined, we can use CELL_INDICES_RBL to find the coset factor of a
+ * cell. In particular, for cell i, its coset factor is roots_of_unity[CELLS_INDICES_RBL[i]]
  */
 static const uint64_t CELL_INDICES_RBL[CELLS_PER_EXT_BLOB] = {
     0x00, 0x40, 0x20, 0x60, 0x10, 0x50, 0x30, 0x70, 0x08, 0x48, 0x28, 0x68, 0x18, 0x58, 0x38, 0x78,
@@ -530,6 +533,72 @@ static C_KZG_RET compute_weighted_sum_of_commitments(
     return ret;
 }
 
+/**
+ * Compute the inverse coset factor h_k^{-1},
+*  where `h_k` is the coset factor for cell with index `k`.
+ *
+ * @param[out]  inv_coset_factor_out  Pointer to store the computed inverse coset factor
+ * @param[in]   cell_index            The index of the cell
+ * @param[in]   s                     The trusted setup
+ */
+static void get_inv_coset_shift_for_cell(
+    fr_t *inv_coset_factor_out,
+    uint64_t cell_index,
+    const KZGSettings *s
+) {
+    /*
+     * Get the cell index in reverse-bit order.
+     * This index points to this cell's coset factor h_k in the roots_of_unity array.
+     */
+    uint64_t cell_idx_rbl = CELL_INDICES_RBL[cell_index];
+
+    /*
+     * Observe that for every element in roots_of_unity, we can find its inverse by
+     * accessing its reflected element.
+     *
+     * For example, consider a multiplicative subgroup with eight elements:
+     * roots = {w^0, w^1, w^2, ... w^7, w^0}
+     * For a root of unity in roots[i], we can find its inverse in roots[-i].
+     */
+    uint64_t inv_coset_factor_idx = FIELD_ELEMENTS_PER_EXT_BLOB - cell_idx_rbl;
+
+    /* Get h_k^{-1} using the index */
+    assert(inv_coset_factor_idx < FIELD_ELEMENTS_PER_EXT_BLOB + 1);
+    *inv_coset_factor_out = s->roots_of_unity[inv_coset_factor_idx];
+}
+
+
+/**
+ * Compute h_k^{n}, where `h_k` is the coset factor for cell with index `k`.
+ *
+ * @param[out]  coset_factor_out  Pointer to store h_k^{n}
+ * @param[in]   cell_index        The index of the cell
+ * @param[in]   s                 The trusted setup
+ */
+static void get_coset_shift_pow_for_cell(
+    fr_t *coset_factor_out,
+    uint64_t cell_index,
+    const KZGSettings *s
+) {
+    /*
+     * Get the cell index in reverse-bit order.
+     * This index points to this cell's coset factor h_k in the roots_of_unity array.
+     */
+    uint64_t cell_idx_rbl = CELL_INDICES_RBL[cell_index];
+
+    /*
+     * Get the index to h_k^n in the roots_of_unity array.
+     *
+     * Multiplying the index of h_k by n, effectively raises h_k to the n-th power,
+     * because advancing in the roots_of_unity array corresponds to increasing exponents.
+     */
+    uint64_t h_k_pow_idx = cell_idx_rbl * FIELD_ELEMENTS_PER_CELL;
+
+    /* Get h_k^n using the index */
+    assert(h_k_pow_idx < FIELD_ELEMENTS_PER_EXT_BLOB + 1);
+    *coset_factor_out = s->roots_of_unity[h_k_pow_idx];
+}
+
 /**
  * Aggregate columns, compute the sum of interpolation polynomials, and commit to the result.
  *
@@ -666,12 +735,10 @@ static C_KZG_RET compute_commitment_to_aggregated_interpolation_poly(
         );
         if (ret != C_KZG_OK) goto out;
 
-        /* Calculate index to the inverse root of unity for this cell index */
-        uint64_t inv_coset_factor_idx = -CELL_INDICES_RBL[i] % FIELD_ELEMENTS_PER_EXT_BLOB;
-        /* For readability, assign root to variable using our index */
-        fr_t *inv_coset_factor = &s->roots_of_unity[inv_coset_factor_idx];
-        /* Now divide by the coset shift factor */
-        shift_poly(column_interpolation_poly, FIELD_ELEMENTS_PER_CELL, inv_coset_factor);
+        /* Shift the poly by h_k^{-1} where h_k is the coset factor for this cell */
+        fr_t inv_coset_factor;
+        get_inv_coset_shift_for_cell(&inv_coset_factor, i, s);
+        shift_poly(column_interpolation_poly, FIELD_ELEMENTS_PER_CELL, &inv_coset_factor);
 
         /* Update the aggregated poly */
         for (size_t k = 0; k < FIELD_ELEMENTS_PER_CELL; k++) {
@@ -728,12 +795,12 @@ static C_KZG_RET computed_weighted_sum_of_proofs(
     if (ret != C_KZG_OK) goto out;
 
     for (uint64_t i = 0; i < num_cells; i++) {
-        /* Calculate index to h_k^n; a root to some power is another root */
-        uint64_t h_k_pow_idx = CELL_INDICES_RBL[cell_indices[i]] * FIELD_ELEMENTS_PER_CELL;
-        /* For readability, assign root to variable using our index */
-        fr_t *h_k_pow = &s->roots_of_unity[h_k_pow_idx];
+        /* Get scaling factor h_k^n where h_k is the coset factor for this cell */
+        fr_t h_k_pow;
+        get_coset_shift_pow_for_cell(&h_k_pow, cell_indices[i], s);
+
         /* Scale the power of r by h_k^n */
-        blst_fr_mul(&weighted_powers_of_r[i], &r_powers[i], h_k_pow);
+        blst_fr_mul(&weighted_powers_of_r[i], &r_powers[i], &h_k_pow);
     }
 
     ret = g1_lincomb_fast(weighted_proof_sum_out, proofs_g1, weighted_powers_of_r, num_cells);

src/test/tests.c

@@ -1838,6 +1838,52 @@ static void test_deduplicate_commitments__one_commitment(void) {
     ASSERT_EQUALS(indices[0], 0);
 }
 
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// Tests for coset shift factors
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static void test_shift_factors__succeeds(void) {
+    fr_t expected_inv_coset_factor, computed_inv_coset_factor, h_k;
+    fr_t expected_coset_factor_pow, computed_coset_factor_pow;
+    static uint64_t n = FIELD_ELEMENTS_PER_CELL;
+
+    /* Loop over all cells */
+    for (uint64_t cell_index = 0; cell_index < CELLS_PER_EXT_BLOB; cell_index++) {
+        /* Get the cell index in reverse-bit order */
+        uint64_t cell_idx_rbl = reverse_bits_limited(CELLS_PER_EXT_BLOB, cell_index);
+
+        /* Ensure the index is within bounds */
+        assert(cell_idx_rbl < FIELD_ELEMENTS_PER_EXT_BLOB + 1);
+
+        /* Get h_k for this cell */
+        h_k = s.roots_of_unity[cell_idx_rbl];
+
+        /* First we test get_inv_coset_shift_for_cell() */
+
+        /* Compute the expected inverse coset factor */
+        blst_fr_eucl_inverse(&expected_inv_coset_factor, &h_k);
+
+        /* Call the function we are testing */
+        get_inv_coset_shift_for_cell(&computed_inv_coset_factor, cell_index, &s);
+
+        /* Compare the expected and computed inverse coset factors */
+        bool ok = fr_equal(&expected_inv_coset_factor, &computed_inv_coset_factor);
+        ASSERT_EQUALS(ok, true);
+
+        /* Now we test get_coset_shift_pow_for_cell() */
+
+        /* Compute the expected coset factor h_k^n */
+        fr_pow(&expected_coset_factor_pow, &h_k, n);
+
+        /* Now call the function we are testing */
+        get_coset_shift_pow_for_cell(&computed_coset_factor_pow, cell_index, &s);
+
+        /* Compare the expected and computed inverse coset factors */
+        ok = fr_equal(&expected_coset_factor_pow, &computed_coset_factor_pow);
+        ASSERT_EQUALS(ok, true);
+    }
+}
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // Tests for recover_cells_and_kzg_proofs
 ////////////////////////////////////////////////////////////////////////////////////////////////////
@@ -2313,6 +2359,7 @@ int main(void) {
     RUN(test_deduplicate_commitments__no_commitments);
     RUN(test_deduplicate_commitments__one_commitment);
     RUN(test_recover_cells_and_kzg_proofs__succeeds_random_blob);
+    RUN(test_shift_factors__succeeds);
     RUN(test_compute_vanishing_polynomial_from_roots);
     RUN(test_vanishing_polynomial_for_missing_cells);
     RUN(test_verify_cell_kzg_proof_batch__succeeds_random_blob);