Added test for reduction + basis for merkle trees

Cargo.toml

@@ -3,7 +3,24 @@ name = "fri"
 version = "0.1.0"
 edition = "2021"
 
+[lib]
+
 [dependencies]
 ark-ff = "0.4.2"
 ark-poly = "0.4.2"
+ark-serialize = "0.4.2"
 rs_merkle = "1.4.2"
+blake3 = { version = "1.5.1", optional = true }
+sha3 = { version = "0.10.8", optional = true }
+
+[dev-dependencies]
+winter-fri = "0.9"
+winter-utils = "0.9"
+winter-rand-utils = "0.9"
+winter-math = "0.9"
+
+[features]
+sha3 = ["dep:sha3"]
+sha3_asm = ["sha3", "sha3/asm"]
+blake3 = ["dep:blake3"]
+default = ["sha3", "blake3"]

src/blake3.rs

+use rs_merkle::Hasher;
+
+/// Blake3-256 implementation of the [`rs_merkle::Hasher`] trait.
+///
+/// See documentation of crate [`sha3`].
+#[derive(Clone, Copy, Debug)]
+pub struct Blake3;
+
+impl Hasher for Blake3 {
+    type Hash = [u8; blake3::OUT_LEN];
+
+    fn hash(data: &[u8]) -> Self::Hash {
+        blake3::hash(data).into()
+    }
+}

src/lib.rs

-use ark_ff::FftField;
+use std::{borrow::Cow, fmt::Debug};
+
+use ark_ff::{FftField, Field};
 use ark_poly::{
-    univariate::DensePolynomial, DenseUVPolynomial, EvaluationDomain, GeneralEvaluationDomain,
-    Polynomial,
+    univariate::DensePolynomial, EvaluationDomain, GeneralEvaluationDomain, Polynomial,
 };
+use ark_serialize::Compress;
 use rs_merkle::{Hasher, MerkleTree};
 
-pub fn commit_polynomial<const N: usize, F: FftField, P: DenseUVPolynomial<F>>(
-    polynomial: &P,
-    commitments: &mut Vec<Vec<F>>,
+#[cfg(feature = "blake3")]
+pub mod blake3;
+#[cfg(feature = "sha3")]
+pub mod sha3;
+
+#[cfg(test)]
+pub mod tests;
+
+pub struct FriCommitments<H: Hasher> {
+    layers: Vec<MerkleTree<H>>,
+}
+
+impl<H: Hasher> FriCommitments<H> {
+    pub fn layers(&self) -> &[MerkleTree<H>] {
+        &self.layers
+    }
+}
+
+impl<H: Hasher> Debug for FriCommitments<H>
+where
+    H::Hash: Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut str = f.debug_struct("FriCommitments");
+        str.field("layers", &MerkleTreeDebug(&self.layers)).finish()
+    }
+}
+
+struct MerkleTreeDebug<'a, H: Hasher>(&'a [MerkleTree<H>]);
+
+impl<H: Hasher> Debug for MerkleTreeDebug<'_, H>
+where
+    H::Hash: Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_list()
+            .entries(self.0.iter().map(MerkleTree::root))
+            .finish()
+    }
+}
+
+struct Layer<H: Hasher> {
+    seed: H::Hash,
+    tree: MerkleTree<H>,
+}
+
+impl<H: Hasher> Layer<H> {
+    #[inline]
+    #[must_use]
+    pub fn new<F: Field>(evaluations: &[F]) -> Self {
+        let tree = build_merkle_tree(evaluations);
+        let seed = tree.root().expect("failed to get merkle tree root");
+        Self { seed, tree }
+    }
+}
+
+struct Prover<H: Hasher> {
+    seed: H::Hash,
+    layers: Vec<MerkleTree<H>>,
+}
+
+impl<H: Hasher> Prover<H> {
+    #[inline]
+    #[must_use]
+    pub fn new<F: Field>(evaluations: &[F], folding_factor: usize) -> Self {
+        let layer = Layer::new(evaluations);
+
+        let mut layers =
+            Vec::with_capacity(evaluations.len().ilog2().div_ceil(folding_factor.ilog2()) as usize);
+        layers.push(layer.tree);
+        Self {
+            seed: layer.seed,
+            layers,
+        }
+    }
+    #[inline]
+    pub fn commit<F: Field>(&mut self, evaluations: &[F]) {
+        let layer = Layer::new(evaluations);
+        self.seed = layer.seed;
+        self.layers.push(layer.tree);
+    }
+    #[inline]
+    pub fn draw_alpha<F: FftField>(&mut self) -> F
+    where
+        H::Hash: AsRef<[u8]>,
+    {
+        for i in 0u64..1000 {
+            let hash = H::concat_and_hash(&self.seed, Some(&H::hash(&i.to_le_bytes())));
+
+            if let Some(elt) = F::from_random_bytes(hash.as_ref()) {
+                return elt;
+            }
+        }
+        panic!("Failed to draw alpha after 1000 attempts");
+    }
+    #[inline]
+    #[must_use]
+    pub fn finish(self) -> FriCommitments<H> {
+        FriCommitments {
+            layers: self.layers,
+        }
+    }
+}
+
+/// Commits the polynomial according to FRI algorithm.
+///
+/// - `polynomial` is the list of coefficients of the polynomial
+/// 
+/// TODO write documentation
+#[must_use]
+pub fn commit_polynomial<const N: usize, F, P, H>(
+    polynomial: Vec<F>,
     blowup_factor: usize,
     remainder_degree: usize,
-) {
-    let domain_size = (polynomial.coeffs().len() * blowup_factor)
+) -> FriCommitments<H>
+where
+    F: FftField,
+    H: Hasher,
+    H::Hash: AsRef<[u8]>,
+{
+    // Convert the polynonial to evaluation form:
+    let domain_size = (polynomial.len() * blowup_factor)
         .checked_next_power_of_two()
-        .expect(&format!(
-            "Domain size out of bounds for blowup factor {blowup_factor} and polynomial of degree-bound {}", polynomial.coeffs().len()
+        .unwrap_or_else(|| panic!(
+            "Domain size out of bounds for blowup factor {blowup_factor} and polynomial of degree-bound {}", polynomial.len()
         ));
+    let mut evaluations = to_evaluations(polynomial, domain_size);
 
-    let domain = GeneralEvaluationDomain::<F>::new(domain_size).unwrap();
-    let mut evaluations = polynomial.coeffs().to_vec();
-    domain.fft_in_place(&mut evaluations);
-
+    // Reduce the polynomial from its evaluations:
+    let mut prover = Prover::new(&evaluations, N);
     let domain = GeneralEvaluationDomain::<F>::new(N).unwrap();
     while evaluations.len() > remainder_degree {
-        commitments.push(evaluations.clone());
-        evaluations = reduce_polynomial::<N, _, _>(&evaluations, &domain, F::GENERATOR);
+        evaluations = reduce_polynomial::<N, _>(&evaluations, prover.draw_alpha(), Some(&domain));
+        prover.commit(&evaluations);
     }
-
-    // TODO commit last
-    commitments.push(evaluations);
+    prover.finish()
 }
 
-fn reduce_polynomial<const N: usize, F: FftField, D: EvaluationDomain<F>>(
+/// Reduces the polynomial by factor `N` using FRI algorithm.
+///
+/// - `N` is the reduction factor. It must be a power of two. Typical values include 2, 4 and 8.
+/// - `evaluations` is the evaluations of the polynomial on the `n`^th roots of unity, where `n` is the
+///    len of `evaluations`. `n` must be a power of two strictly greater than the degree-bound of the polynomial.
+///    If `w` is `F::get_root_of_unity(n).unwrap()`, `evaluations[i]` is the evaluation at `w^i`.
+/// - `alpha` is the "challenge" used to reduce the polynomial.
+/// - `domain`, if provided, is the pre-computed evaluation domain of size `N`.
+///
+/// # Returns
+/// If `evaluations` corresponds to `P(X) = a_0 + a_1 X + ... + a_(n-1) X^(n-1)` in coefficient form, then `P` is
+/// decomposed in `P_i(X) = a_i + a_(N+i) X + ... + a_(kN+i) X^k` where `k=n/N`.
+///
+/// This function returns the evaluations of `Q(X) = P_0(X^N) + alpha P_1(X^N) + ... + alpha^(N-1) P_(N-1)(X^N)` on
+/// the `n/N`th roots of unity.
+///
+/// # Panics
+/// This may panic if `N` or `evaluations.len()` are not powers of two, if `N > evaluations.len()` and if
+/// `F` does not contain subgroups of size `evaluations.len()` and `N`.
+///
+/// # Credits
+/// This is partly based on equation (4) from [https://eprint.iacr.org/2022/1216.pdf].
+#[must_use]
+pub fn reduce_polynomial<const N: usize, F: FftField>(
     evaluations: &[F],
-    domain: &D,
     alpha: F,
+    domain: Option<&GeneralEvaluationDomain<F>>,
 ) -> Vec<F> {
+    let domain = domain.map_or_else(
+        || Cow::Owned(GeneralEvaluationDomain::new(N).unwrap()),
+        Cow::Borrowed,
+    );
+
     debug_assert!(
         evaluations.len().is_power_of_two(),
         "Number of evaluations must be a power of two"
@@ -58,12 +200,13 @@ fn reduce_polynomial<const N: usize, F: FftField, D: EvaluationDomain<F>>(
     let mut offset = F::ONE;
 
     for i in 0..bound {
-        buffer.extend(evaluations.iter().skip(i).step_by(bound));
+        buffer.extend(evaluations.iter().skip(i).step_by(bound).copied());
         domain.ifft_in_place(&mut buffer);
 
         let mut factor = F::ONE;
         for coeff in &mut buffer {
-            *coeff *= factor;
+            // FIXME: rust-analyzer fails to infer type of `coeff` on VS Code
+            *(coeff as &mut F) *= factor;
             factor *= offset;
         }
         offset *= root_inv;
@@ -76,3 +219,29 @@ fn reduce_polynomial<const N: usize, F: FftField, D: EvaluationDomain<F>>(
     }
     new_evaluations
 }
+
+#[inline]
+pub fn to_evaluations<F: FftField>(mut polynomial: Vec<F>, domain_size: usize) -> Vec<F> {
+    debug_assert!(
+        domain_size.is_power_of_two(),
+        "Domain size must be a power of two"
+    );
+
+    let domain = GeneralEvaluationDomain::<F>::new(domain_size).unwrap();
+    domain.fft_in_place(&mut polynomial);
+    polynomial
+}
+
+#[must_use]
+fn build_merkle_tree<F: Field, H: Hasher>(evaluations: &[F]) -> MerkleTree<H> {
+    let mut hashes = Vec::with_capacity(evaluations.len());
+    let mut bytes = vec![];
+    for evaluation in evaluations {
+        evaluation
+            .serialize_with_mode(&mut bytes, Compress::Yes)
+            .expect("Serialization failed");
+        hashes.push(H::hash(&bytes));
+        bytes.clear();
+    }
+    MerkleTree::from_leaves(&hashes)
+}

src/sha3.rs

+use rs_merkle::Hasher;
+use sha3::{digest::FixedOutput, Digest};
+
+/// Sha3-256 implementation of the [`rs_merkle::Hasher`] trait.
+/// 
+/// See documentation of crate [`sha3`].
+#[derive(Clone, Copy, Debug)]
+pub struct Sha3_256;
+
+impl Hasher for Sha3_256 {
+    type Hash = [u8; 32];
+
+    fn hash(data: &[u8]) -> Self::Hash {
+        let mut hasher = sha3::Sha3_256::new();
+        hasher.update(data);
+        <[u8; 32]>::from(hasher.finalize_fixed())
+    }
+}
+
+/// Sha3-512 implementation of the [`rs_merkle::Hasher`] trait.
+/// 
+/// See documentation of crate [`sha3`].
+#[derive(Clone, Copy, Debug)]
+pub struct Sha3_512;
+
+impl Hasher for Sha3_512 {
+    type Hash = [u8; 64];
+
+    fn hash(data: &[u8]) -> Self::Hash {
+        let mut hasher = sha3::Sha3_512::new();
+        hasher.update(data);
+        <[u8; 64]>::from(hasher.finalize_fixed())
+    }
+}
\ No newline at end of file

src/tests.rs

+use ark_ff::{Fp128, MontBackend, MontConfig};
+use winter_math::{fields::f128::BaseElement, FieldElement, StarkField};
+use winter_rand_utils::{rand_value, rand_vector};
+use winter_utils::transpose_slice;
+
+use crate::to_evaluations;
+
+const NUMBER_OF_POLYNOMIALS: usize = 10;
+const POLY_COEFFS_LEN: usize = 2048;
+const BLOWUP_FACTOR: usize = 4;
+
+const DOMAIN_SIZE: usize = (POLY_COEFFS_LEN * BLOWUP_FACTOR).next_power_of_two();
+
+/// Matches `BaseElement` from winterfell
+#[derive(MontConfig)]
+#[modulus = "340282366920938463463374557953744961537"]
+#[generator = "3"]
+pub struct Test;
+/// A prime, fft-friendly field isomorph to [`winter_math::fields::f128::BaseElement`].
+pub type Fq = Fp128<MontBackend<Test, 2>>;
+
+macro_rules! do_for_multiple_folding_factors {
+    ($factor: ident = $($factors: literal),* => $action: block) => {
+        {
+            $({
+                const $factor: usize = $factors;
+                $action;
+            })*
+        }
+    };
+}
+
+// This assumes winterfri is correct
+#[test]
+fn test_reduction() {
+    for i in 0..NUMBER_OF_POLYNOMIALS {
+        println!("Testing on polynomial {i}");
+
+        let poly = rand_vector(DOMAIN_SIZE);
+        let poly2 = convert_many(&poly);
+
+        do_for_multiple_folding_factors!(FACTOR = 2, 4, 8, 16 => {
+            println!("--Folding factor={FACTOR}");
+            let mut evaluations = prepare_winterfell_poly(poly.clone());
+            let mut evaluations2 = to_evaluations(poly2.clone(), DOMAIN_SIZE);
+            assert_eq!(convert_many(&evaluations), evaluations2);
+
+            while evaluations.len() > FACTOR {
+                let alpha = rand_value();
+                let alpha2 = convert(&alpha);
+
+                let transposed = transpose_slice::<_, FACTOR>(&evaluations);
+                evaluations = winter_fri::folding::apply_drp(&transposed, BaseElement::ONE, alpha);
+                evaluations2 = super::reduce_polynomial::<FACTOR, _>(&evaluations2, alpha2, None);
+
+                assert_eq!(convert_many(&evaluations), evaluations2);
+            }
+        });
+    }
+}
+
+#[inline]
+fn prepare_winterfell_poly<F: StarkField>(mut poly: Vec<F>) -> Vec<F> {
+    use winter_math::fft::{evaluate_poly, get_twiddles};
+    poly.resize(DOMAIN_SIZE, F::ZERO);
+    let twiddles = get_twiddles(DOMAIN_SIZE);
+    evaluate_poly(&mut poly, &twiddles);
+    poly
+}
+
+#[inline]
+pub fn convert(value: &BaseElement) -> Fq {
+    Fq::from(value.as_int())
+}
+
+#[inline]
+pub fn convert_many(values: &[BaseElement]) -> Vec<Fq> {
+    values.iter().map(convert).collect()
+}