Fri refactor (Plonky3#224)

* initial fri refactor

* progress

* more work

* hmm debugging

* fix bug

* well it works it's just slow

* delete p3-ldt

* rm p3-ldt

* do math smarter

* wee faster

* enable parallel

* most optimization done, time to start cleaning up

* linting

* verifier works

* cleanup

* lints and ci

* fmt

* no_std

* rm commented code

* fix error handling

* iterate theh right way

* hm, it works even with a normal Vec

* initial config rework

* make config more concise

* fix tests

* clippy

* remove commented code

* move two_adic_fri_pcs to toplevel reexport

* fmt

Cargo.toml

@@ -17,7 +17,6 @@ members = [
     "keccak",
     "keccak-air",
     "lde",
-    "ldt",
     "matrix",
     "merkle-tree",
     "maybe-rayon",

challenger/src/duplex_challenger.rs

@@ -1,7 +1,7 @@
 use alloc::vec;
 use alloc::vec::Vec;
 
-use p3_field::PrimeField64;
+use p3_field::{ExtensionField, Field, PrimeField64};
 use p3_symmetric::CryptographicPermutation;
 
 use crate::{CanObserve, CanSample, CanSampleBits, FieldChallenger};
@@ -87,21 +87,24 @@ where
     }
 }
 
-impl<F, P, const WIDTH: usize> CanSample<F> for DuplexChallenger<F, P, WIDTH>
+impl<F, EF, P, const WIDTH: usize> CanSample<EF> for DuplexChallenger<F, P, WIDTH>
 where
-    F: Copy,
+    F: Field,
+    EF: ExtensionField<F>,
     P: CryptographicPermutation<[F; WIDTH]>,
 {
-    fn sample(&mut self) -> F {
-        // If we have buffered inputs, we must perform a duplexing so that the challenge will
-        // reflect them. Or if we've run out of outputs, we must perform a duplexing to get more.
-        if !self.input_buffer.is_empty() || self.output_buffer.is_empty() {
-            self.duplexing();
-        }
-
-        self.output_buffer
-            .pop()
-            .expect("Output buffer should be non-empty")
+    fn sample(&mut self) -> EF {
+        EF::from_base_fn(|_| {
+            // If we have buffered inputs, we must perform a duplexing so that the challenge will
+            // reflect them. Or if we've run out of outputs, we must perform a duplexing to get more.
+            if !self.input_buffer.is_empty() || self.output_buffer.is_empty() {
+                self.duplexing();
+            }
+
+            self.output_buffer
+                .pop()
+                .expect("Output buffer should be non-empty")
+        })
     }
 }
 
@@ -218,7 +221,9 @@ mod tests {
 
         (0..WIDTH / 2).for_each(|element| {
             assert_eq!(
-                duplex_challenger.sample(),
+                <DuplexChallenger<F, TestPermutation, WIDTH> as CanSample<F>>::sample(
+                    &mut duplex_challenger
+                ),
                 F::from_canonical_u8(element as u8)
             );
             assert_eq!(
@@ -229,7 +234,12 @@ mod tests {
         });
 
         (0..WIDTH / 2).for_each(|i| {
-            assert_eq!(duplex_challenger.sample(), F::from_canonical_u8(0));
+            assert_eq!(
+                <DuplexChallenger<F, TestPermutation, WIDTH> as CanSample<F>>::sample(
+                    &mut duplex_challenger
+                ),
+                F::from_canonical_u8(0)
+            );
             assert_eq!(duplex_challenger.input_buffer, vec![]);
             assert_eq!(
                 duplex_challenger.output_buffer,

challenger/src/grinding_challenger.rs

@@ -1,34 +1,39 @@
-use p3_field::PrimeField64;
+use p3_field::{Field, PrimeField64};
 use p3_maybe_rayon::prelude::*;
 use p3_symmetric::CryptographicPermutation;
 use tracing::instrument;
 
-use crate::{DuplexChallenger, FieldChallenger};
+use crate::{CanObserve, CanSampleBits, DuplexChallenger};
 
-pub trait GrindingChallenger<F: PrimeField64>: FieldChallenger<F> + Clone {
-    // Can be overridden for more efficient methods not involving cloning, depending on the
-    // internals of the challenger.
-    #[instrument(name = "grind for proof-of-work witness", skip_all)]
-    fn grind(&mut self, bits: usize) -> F {
-        let witness = (0..F::ORDER_U64)
-            .into_par_iter()
-            .map(|i| F::from_canonical_u64(i))
-            .find_any(|witness| self.clone().check_witness(bits, *witness))
-            .expect("failed to find witness");
-        assert!(self.check_witness(bits, witness));
-        witness
-    }
+pub trait GrindingChallenger:
+    CanObserve<Self::Witness> + CanSampleBits<usize> + Sync + Clone
+{
+    type Witness: Field;
+
+    fn grind(&mut self, bits: usize) -> Self::Witness;
 
     #[must_use]
-    fn check_witness(&mut self, bits: usize, witness: F) -> bool {
+    fn check_witness(&mut self, bits: usize, witness: Self::Witness) -> bool {
         self.observe(witness);
         self.sample_bits(bits) == 0
     }
 }
 
-impl<F, P, const WIDTH: usize> GrindingChallenger<F> for DuplexChallenger<F, P, WIDTH>
+impl<F, P, const WIDTH: usize> GrindingChallenger for DuplexChallenger<F, P, WIDTH>
 where
     F: PrimeField64,
     P: CryptographicPermutation<[F; WIDTH]>,
 {
+    type Witness = F;
+
+    #[instrument(name = "grind for proof-of-work witness", skip_all)]
+    fn grind(&mut self, bits: usize) -> Self::Witness {
+        let witness = (0..F::ORDER_U64)
+            .into_par_iter()
+            .map(|i| F::from_canonical_u64(i))
+            .find_any(|witness| self.clone().check_witness(bits, *witness))
+            .expect("failed to find witness");
+        assert!(self.check_witness(bits, witness));
+        witness
+    }
 }

fri/Cargo.toml

@@ -7,8 +7,9 @@ license = "MIT OR Apache-2.0"
 [dependencies]
 p3-challenger = { path = "../challenger" }
 p3-commit = { path = "../commit" }
+p3-dft = { path = "../dft" }
 p3-field = { path = "../field" }
-p3-ldt = { path = "../ldt" }
+p3-interpolation = { path = "../interpolation" }
 p3-matrix = { path = "../matrix" }
 p3-maybe-rayon = { path = "../maybe-rayon" }
 p3-util = { path = "../util" }

fri/src/config.rs

@@ -1,88 +1,12 @@
-use core::marker::PhantomData;
-
-use p3_challenger::{CanObserve, GrindingChallenger};
-use p3_commit::{DirectMmcs, Mmcs};
-use p3_field::{ExtensionField, PrimeField64, TwoAdicField};
-
-pub trait FriConfig {
-    type Val: PrimeField64;
-    type Challenge: ExtensionField<Self::Val> + TwoAdicField;
-
-    type InputMmcs: Mmcs<Self::Val>;
-    type CommitPhaseMmcs: DirectMmcs<Self::Challenge>;
-
-    type Challenger: GrindingChallenger<Self::Val>
-        + CanObserve<<Self::CommitPhaseMmcs as Mmcs<Self::Challenge>>::Commitment>;
-
-    fn commit_phase_mmcs(&self) -> &Self::CommitPhaseMmcs;
-
-    fn num_queries(&self) -> usize;
-
-    fn log_blowup(&self) -> usize;
-
-    fn blowup(&self) -> usize {
-        1 << self.log_blowup()
-    }
-
-    fn proof_of_work_bits(&self) -> usize;
-}
-
-pub struct FriConfigImpl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger> {
-    log_blowup: usize,
-    num_queries: usize,
-    proof_of_work_bits: usize,
-    commit_phase_mmcs: CommitPhaseMmcs,
-    _phantom: PhantomData<(Val, Challenge, InputMmcs, Challenger)>,
-}
-
-impl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger>
-    FriConfigImpl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger>
-{
-    pub fn new(
-        log_blowup: usize,
-        num_queries: usize,
-        proof_of_work_bits: usize,
-        commit_phase_mmcs: CommitPhaseMmcs,
-    ) -> Self {
-        Self {
-            log_blowup,
-            num_queries,
-            commit_phase_mmcs,
-            proof_of_work_bits,
-            _phantom: PhantomData,
-        }
-    }
+pub struct FriConfig<M> {
+    pub log_blowup: usize,
+    pub num_queries: usize,
+    pub proof_of_work_bits: usize,
+    pub mmcs: M,
 }
 
-impl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger> FriConfig
-    for FriConfigImpl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger>
-where
-    Val: PrimeField64,
-    Challenge: ExtensionField<Val> + TwoAdicField,
-    InputMmcs: Mmcs<Val>,
-    CommitPhaseMmcs: DirectMmcs<Challenge>,
-    Challenger:
-        GrindingChallenger<Val> + CanObserve<<CommitPhaseMmcs as Mmcs<Challenge>>::Commitment>,
-{
-    type Val = Val;
-    type Challenge = Challenge;
-    type InputMmcs = InputMmcs;
-    type CommitPhaseMmcs = CommitPhaseMmcs;
-    type Challenger = Challenger;
-
-    fn commit_phase_mmcs(&self) -> &CommitPhaseMmcs {
-        &self.commit_phase_mmcs
-    }
-
-    fn num_queries(&self) -> usize {
-        self.num_queries
-    }
-
-    fn log_blowup(&self) -> usize {
-        self.log_blowup
-    }
-
-    fn proof_of_work_bits(&self) -> usize {
-        self.proof_of_work_bits
+impl<M> FriConfig<M> {
+    pub fn blowup(&self) -> usize {
+        1 << self.log_blowup
     }
 }

fri/src/lib.rs

@@ -4,72 +4,14 @@
 
 extern crate alloc;
 
-use alloc::vec::Vec;
-
-use p3_commit::Mmcs;
-use p3_ldt::{Ldt, LdtBasedPcs};
-use p3_matrix::Dimensions;
-use verifier::VerificationErrorForFriConfig;
-
-use crate::prover::prove;
-use crate::verifier::verify;
-
 mod config;
 mod fold_even_odd;
-mod matrix_reducer;
 mod proof;
-mod prover;
-mod verifier;
+pub mod prover;
+pub mod two_adic_pcs;
+pub mod verifier;
 
 pub use config::*;
 pub use fold_even_odd::*;
 pub use proof::*;
-
-pub struct FriLdt<FC: FriConfig> {
-    pub config: FC,
-}
-
-impl<FC: FriConfig> Ldt<FC::Val, FC::InputMmcs, FC::Challenger> for FriLdt<FC> {
-    type Proof = FriProof<FC>;
-    type Error = VerificationErrorForFriConfig<FC>;
-
-    fn log_blowup(&self) -> usize {
-        self.config.log_blowup()
-    }
-
-    fn prove(
-        &self,
-        input_mmcs: &[FC::InputMmcs],
-        input_data: &[&<FC::InputMmcs as Mmcs<FC::Val>>::ProverData],
-        challenger: &mut FC::Challenger,
-    ) -> Self::Proof {
-        prove::<FC>(&self.config, input_mmcs, input_data, challenger)
-    }
-
-    fn verify(
-        &self,
-        input_mmcs: &[FC::InputMmcs],
-        input_dims: &[Vec<Dimensions>],
-        input_commits: &[<FC::InputMmcs as Mmcs<FC::Val>>::Commitment],
-        proof: &Self::Proof,
-        challenger: &mut FC::Challenger,
-    ) -> Result<(), Self::Error> {
-        verify::<FC>(
-            &self.config,
-            input_mmcs,
-            input_dims,
-            input_commits,
-            proof,
-            challenger,
-        )
-    }
-}
-
-pub type FriBasedPcs<FC, Mmcs, Dft, Challenger> = LdtBasedPcs<
-    <FC as FriConfig>::Val,
-    <FC as FriConfig>::Challenge,
-    Dft,
-    Mmcs,
-    FriLdt<FC>,
-    Challenger,
->;
+pub use two_adic_pcs::*;