From 10bff7d5c8f202cf40f275757e8af2f8c50dc7b9 Mon Sep 17 00:00:00 2001
From: Ramana Kumar <ramana@member.fsf.org>
Date: Sat, 17 Sep 2022 10:49:04 +0100
Subject: [PATCH] Fix eval_poly_l for evaluation at a root

Move the TODO for the special formula to compute_proof_single_l
---
 src/kzg_proofs.c | 36 ++++++++++++++++++++++++++++++++++++
 src/poly.c       |  6 ++++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/kzg_proofs.c b/src/kzg_proofs.c
index aff1fe1..90fbb83 100644
--- a/src/kzg_proofs.c
+++ b/src/kzg_proofs.c
@@ -92,6 +92,7 @@ C_KZG_RET check_proof_single(bool *out, const g1_t *commitment, const g1_t *proo
 }
 
 // TODO: I don't think this should compute the evaluation. Instead y should be a parameter
+// TODO: Consider the case where x is one of the roots of unity (needs special formula)
 C_KZG_RET compute_proof_single_l(g1_t *out, const poly_l *p, const fr_t *x0, const KZGSettings *ks) {
   fr_t y, tmp, tmp2;
   poly_l q;
@@ -550,9 +551,44 @@ void poly_eval_l_check(void) {
     free_poly_l(&p_l);
 }
 
+void eval_poly_l_at_first_root_of_unity(void) {
+    uint64_t n = 10;
+    fr_t actual, expected;
+    poly p;
+    new_poly(&p, n);
+    for (uint64_t i = 0; i < n; i++) {
+        fr_from_uint64(&p.coeffs[i], i + 2);
+    }
+
+    poly_l p_l;
+    FFTSettings fs;
+    KZGSettings ks;
+    uint64_t secrets_len = 16;
+    g1_t s1[secrets_len];
+    g2_t s2[secrets_len];
+
+    generate_trusted_setup(s1, s2, &secret, secrets_len);
+    TEST_CHECK(C_KZG_OK == new_fft_settings(&fs, 4)); // log_2(secrets_len)
+    TEST_CHECK(C_KZG_OK == new_kzg_settings(&ks, s1, s2, secrets_len, &fs));
+
+    eval_poly(&expected, &p, &fs.expanded_roots_of_unity[0]);
+
+    TEST_CHECK(C_KZG_OK == new_poly_l_from_poly(&p_l, &p, &ks));
+
+    eval_poly_l(&actual, &p_l, &fs.expanded_roots_of_unity[0], &fs);
+
+    TEST_CHECK(fr_equal(&expected, &actual));
+
+    free_fft_settings(&fs);
+    free_kzg_settings(&ks);
+    free_poly(&p);
+    free_poly_l(&p_l);
+}
+
 TEST_LIST = {
     {"KZG_PROOFS_TEST", title},
     {"poly_eval_l_check", poly_eval_l_check},
+    {"eval_poly_l_at_first_root_of_unity", eval_poly_l_at_first_root_of_unity},
     {"proof_single", proof_single},
     {"proof_multi", proof_multi},
     {"commit_to_nil_poly", commit_to_nil_poly},
diff --git a/src/poly.c b/src/poly.c
index aebeb9f..34346ed 100644
--- a/src/poly.c
+++ b/src/poly.c
@@ -103,8 +103,6 @@ void eval_poly(fr_t *out, const poly *p, const fr_t *x) {
 }
 
 // TODO: optimize via batch inversion
-// TODO: Consider the case where x is one of the roots of unity
-//       (needs special formula)
 void eval_poly_l(fr_t *out, const poly_l *p, const fr_t *x, const FFTSettings *fs) {
   fr_t tmp;
   uint64_t i;
@@ -112,6 +110,10 @@ void eval_poly_l(fr_t *out, const poly_l *p, const fr_t *x, const FFTSettings *f
 
   *out = fr_zero;
   for (i = 0; i < p->length; i++) {
+    if (fr_equal(x, &fs->expanded_roots_of_unity[i * stride])) {
+      *out = p->values[i];
+      return;
+    }
     fr_sub(&tmp, x, &fs->expanded_roots_of_unity[i * stride]);
     fr_inv(&tmp, &tmp);
     fr_mul(&tmp, &tmp, &fs->expanded_roots_of_unity[i * stride]);