From 0fa0942981124d87a5d33cf78e75f02936e7854c Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Fri, 30 Apr 2021 15:07:54 +0200
Subject: [PATCH] FRI on coset

---
 src/fri.rs                   | 13 +++++++----
 src/polynomial/commitment.rs | 26 +++++++++++++++++++++
 src/polynomial/mod.rs        |  1 +
 src/polynomial/polynomial.rs | 45 ++++++++++++++++++++++++++++++++++++
 4 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 src/polynomial/commitment.rs

diff --git a/src/fri.rs b/src/fri.rs
index 9376f449..849cebc0 100644
--- a/src/fri.rs
+++ b/src/fri.rs
@@ -98,6 +98,7 @@ fn fri_committed_trees<F: Field>(
 
     let mut trees = Vec::new();
 
+    let mut shift = F::MULTIPLICATIVE_GROUP_GENERATOR;
     let num_reductions = config.reduction_arity_bits.len();
     for i in 0..num_reductions {
         let arity = 1 << config.reduction_arity_bits[i];
@@ -124,8 +125,9 @@ fn fri_committed_trees<F: Field>(
                 .map(|chunk| reduce_with_powers(chunk, beta))
                 .collect::<Vec<_>>(),
         );
+        shift = shift.exp_u32(arity as u32);
         // TODO: Is it faster to interpolate?
-        values = fft(coeffs.clone());
+        values = coeffs.clone().coset_fft(shift);
     }
 
     challenger.observe_elements(&coeffs.coeffs);
@@ -315,7 +317,8 @@ fn fri_verifier_query_round<F: Field>(
     let mut old_x_index = 0;
     // `subgroup_x` is `subgroup[x_index]`, i.e., the actual field element in the domain.
     let log_n = log2_strict(n);
-    let mut subgroup_x = F::primitive_root_of_unity(log_n).exp_usize(reverse_bits(x_index, log_n));
+    let mut subgroup_x = F::MULTIPLICATIVE_GROUP_GENERATOR
+        * F::primitive_root_of_unity(log_n).exp_usize(reverse_bits(x_index, log_n));
     for (i, &arity_bits) in config.reduction_arity_bits.iter().enumerate() {
         let arity = 1 << arity_bits;
         let next_domain_size = domain_size >> arity_bits;
@@ -397,8 +400,8 @@ mod tests {
         type F = CrandallField;
 
         let n = 1 << degree_log;
-        let evals = PolynomialValues::new((0..n).map(|_| F::rand()).collect());
-        let lde = evals.clone().lde(rate_bits);
+        let coeffs = PolynomialCoeffs::new((0..n).map(|_| F::rand()).collect()).lde(rate_bits);
+        let coset_lde = coeffs.clone().coset_fft(F::MULTIPLICATIVE_GROUP_GENERATOR);
         let config = FriConfig {
             num_query_rounds,
             rate_bits,
@@ -406,7 +409,7 @@ mod tests {
             reduction_arity_bits,
         };
         let mut challenger = Challenger::new();
-        let proof = fri_proof(&ifft(lde.clone()), &lde, &mut challenger, &config);
+        let proof = fri_proof(&coeffs, &coset_lde, &mut challenger, &config);
 
         let mut challenger = Challenger::new();
         verify_fri_proof(degree_log, &proof, &mut challenger, &config)?;
diff --git a/src/polynomial/commitment.rs b/src/polynomial/commitment.rs
new file mode 100644
index 00000000..578cc237
--- /dev/null
+++ b/src/polynomial/commitment.rs
@@ -0,0 +1,26 @@
+use crate::field::field::Field;
+use crate::merkle_tree::MerkleTree;
+use crate::polynomial::polynomial::PolynomialValues;
+use crate::util::transpose;
+
+struct ListPolynomialCommitment<F: Field> {
+    pub lde_values: Vec<Vec<F>>,
+    pub rate_bits: usize,
+    pub merkle_tree: MerkleTree<F>,
+}
+
+impl<F: Field> ListPolynomialCommitment<F> {
+    pub fn new(values: Vec<PolynomialValues<F>>, rate_bits: usize) -> Self {
+        let lde_values = values
+            .into_iter()
+            .map(|p| p.lde(rate_bits).values)
+            .collect::<Vec<_>>();
+        let merkle_tree = MerkleTree::new(transpose(&lde_values), false);
+
+        Self {
+            lde_values,
+            rate_bits,
+            merkle_tree,
+        }
+    }
+}
diff --git a/src/polynomial/mod.rs b/src/polynomial/mod.rs
index 2c7f7076..4360d3d4 100644
--- a/src/polynomial/mod.rs
+++ b/src/polynomial/mod.rs
@@ -1,2 +1,3 @@
+pub mod commitment;
 pub(crate) mod division;
 pub mod polynomial;
diff --git a/src/polynomial/polynomial.rs b/src/polynomial/polynomial.rs
index 7034a8b9..07f6434d 100644
--- a/src/polynomial/polynomial.rs
+++ b/src/polynomial/polynomial.rs
@@ -108,4 +108,49 @@ impl<F: Field> PolynomialCoeffs<F> {
             .find(|&i| self.coeffs[i].is_nonzero())
             .map_or(0, |i| i + 1)
     }
+
+    pub fn fft(self) -> PolynomialValues<F> {
+        fft(self)
+    }
+
+    pub fn coset_fft(self, shift: F) -> PolynomialValues<F> {
+        let modified_poly: Self = shift
+            .powers()
+            .zip(self.coeffs)
+            .map(|(r, c)| r * c)
+            .collect::<Vec<_>>()
+            .into();
+        modified_poly.fft()
+    }
+}
+
+impl<F: Field> From<Vec<F>> for PolynomialCoeffs<F> {
+    fn from(coeffs: Vec<F>) -> Self {
+        Self::new(coeffs)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::field::crandall_field::CrandallField;
+
+    #[test]
+    fn test_coset_fft() {
+        type F = CrandallField;
+
+        let k = 8;
+        let n = 1 << k;
+        let poly = PolynomialCoeffs::new((0..n).map(|_| F::rand()).collect());
+        let shift = F::rand();
+        let coset_evals = poly.clone().coset_fft(shift).values;
+
+        let generator = F::primitive_root_of_unity(k);
+        let naive_coset_evals = F::cyclic_subgroup_coset_known_order(generator, shift, n)
+            .into_iter()
+            .map(|x| poly.eval(x))
+            .collect::<Vec<_>>();
+
+        assert_eq!(coset_evals, naive_coset_evals);
+    }
 }