Fix compute_kzg_proof_impl() when z is inside the domain (#111)

* Fix compute_kzg_proof_impl() when `z` is inside the domain

* Satisfy linter D:

src/c_kzg_4844.c

@@ -1103,7 +1103,9 @@ C_KZG_RET compute_kzg_proof_impl(KZGProof *out, const Polynomial *polynomial, co
 
     for (i = 0; i < FIELD_ELEMENTS_PER_BLOB; i++) {
         if (fr_equal(z, &roots_of_unity[i])) {
+            /* We are asked to compute a KZG proof inside the domain */
             m = i + 1;
+            inverses_in[i] = FR_ONE;
             continue;
         }
         // (p_i - y) / (ω_i - z)
@@ -1122,15 +1124,20 @@ C_KZG_RET compute_kzg_proof_impl(KZGProof *out, const Polynomial *polynomial, co
         q.evals[--m] = FR_ZERO;
         for (i = 0; i < FIELD_ELEMENTS_PER_BLOB; i++) {
             if (i == m) continue;
-            // (p_i - y) * ω_i / (z * (z - ω_i))
+            /* Build denominator: z * (z - ω_i) */
             blst_fr_sub(&tmp, z, &roots_of_unity[i]);
             blst_fr_mul(&inverses_in[i], &tmp, z);
         }
+
         ret = fr_batch_inv(inverses, inverses_in, FIELD_ELEMENTS_PER_BLOB);
         if (ret != C_KZG_OK) goto out;
+
         for (i = 0; i < FIELD_ELEMENTS_PER_BLOB; i++) {
+            if (i == m) continue;
+            /* Build numerator: ω_i * (p_i - y) */
             blst_fr_sub(&tmp, &polynomial->evals[i], &y);
             blst_fr_mul(&tmp, &tmp, &roots_of_unity[i]);
+            /* Do the division: (p_i - y) * ω_i / (z * (z - ω_i)) */
             blst_fr_mul(&tmp, &tmp, &inverses[i]);
             blst_fr_add(&q.evals[m], &q.evals[m], &tmp);
         }

src/test_c_kzg_4844.c

@@ -421,7 +421,8 @@ static void test_log_2_byte__expected_values(void) {
     while (true) {
         /*
          * Corresponds to the index of the highest bit set in the byte.
-         * Adapted from https://graphics.stanford.edu/~seander/bithacks.html#IntegerLog.
+         * Adapted from
+         * https://graphics.stanford.edu/~seander/bithacks.html#IntegerLog.
          */
         byte b = i;
         int r, shift;
@@ -442,7 +443,7 @@ static void test_log_2_byte__expected_values(void) {
 // Tests for compute_kzg_proof
 ///////////////////////////////////////////////////////////////////////////////
 
-static void test_compute_and_verify_kzg_proof(void) {
+static void test_compute_and_verify_kzg_proof__succeeds_round_trip(void) {
     C_KZG_RET ret;
     Bytes48 proof;
     Bytes32 z, y;
@@ -487,6 +488,49 @@ static void test_compute_and_verify_kzg_proof(void) {
     ASSERT_EQUALS(ok, 1);
 }
 
+static void test_compute_and_verify_kzg_proof__succeeds_within_domain(void) {
+    const int SAMPLES = 25;
+    for (int i = 0; i < SAMPLES; i++) {
+        C_KZG_RET ret;
+        Blob blob;
+        KZGCommitment c;
+        Polynomial poly;
+        Bytes48 proof;
+        Bytes32 z, y;
+        fr_t y_fr, z_fr;
+        bool ok;
+
+        get_rand_blob(&blob);
+
+        ret = blob_to_kzg_commitment(&c, &blob, &s);
+        ASSERT_EQUALS(ret, C_KZG_OK);
+
+        ret = blob_to_polynomial(&poly, &blob);
+        ASSERT_EQUALS(ret, C_KZG_OK);
+
+        z_fr = s.fs->roots_of_unity[i];
+        bytes_from_bls_field(&z, &z_fr);
+
+        /* Compute the proof */
+        ret = compute_kzg_proof(&proof, &blob, &z, &s);
+        ASSERT_EQUALS(ret, C_KZG_OK);
+
+        /* Now evaluate the poly at `z` to learn `y` */
+        ret = evaluate_polynomial_in_evaluation_form(&y_fr, &poly, &z_fr, &s);
+        ASSERT_EQUALS(ret, C_KZG_OK);
+
+        /* Now also get `y` in bytes */
+        bytes_from_bls_field(&y, &y_fr);
+
+        /* Finally verify the proof */
+        ret = verify_kzg_proof(&ok, &c, &z, &y, &proof, &s);
+        ASSERT_EQUALS(ret, C_KZG_OK);
+
+        /* The proof should verify! */
+        ASSERT_EQUALS(ok, 1);
+    }
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // Main logic
 ///////////////////////////////////////////////////////////////////////////////
@@ -532,7 +576,8 @@ int main(void) {
     RUN(test_reverse_bits__some_bits_are_one);
     RUN(test_reverse_bits__all_bits_are_one);
     RUN(test_log_2_byte__expected_values);
-    RUN(test_compute_and_verify_kzg_proof);
+    RUN(test_compute_and_verify_kzg_proof__succeeds_round_trip);
+    RUN(test_compute_and_verify_kzg_proof__succeeds_within_domain);
     teardown();
 
     return TEST_REPORT();