Merge pull request #16 from succinctlabs/tamir/air-api-change

feat: Plonky3 updates

Cargo.toml

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

README.md

@@ -57,6 +57,14 @@ Hashes
 - [x] Monolith
 
 
+## Benchmark
+
+We sometimes use a Keccak AIR to compare Plonky3's performance to other libraries like Plonky2. Several variations are possible here, with different fields and so forth, but here is one example:
+```
+RUST_LOG=info cargo run --example prove_baby_bear_keccak --release --features parallel
+```
+
+
 ## License
 
 Licensed under either of

air/src/air.rs

@@ -22,7 +22,7 @@ pub trait Air<AB: AirBuilder>: BaseAir<AB::F> {
 pub trait AirBuilder: Sized {
     type F: Field;
 
-    type Expr: AbstractField<F = Self::F>
+    type Expr: AbstractField
         + From<Self::F>
         + Add<Self::Var, Output = Self::Expr>
         + Add<Self::F, Output = Self::Expr>
@@ -106,51 +106,40 @@ pub trait AirBuilder: Sized {
         let x = x.into();
         self.assert_zero(x.clone() * (x - Self::Expr::one()));
     }
-
-    fn assert_zero_ext<ExprExt, I>(&mut self, x: I)
-    where
-        ExprExt: AbstractExtensionField<Self::Expr>,
-        I: Into<ExprExt>,
-    {
-        for xb in x.into().as_base_slice().iter().cloned() {
-            self.assert_zero(xb);
-        }
-    }
-
-    fn assert_eq_ext<ExprExt, I1, I2>(&mut self, x: I1, y: I2)
-    where
-        ExprExt: AbstractExtensionField<Self::Expr>,
-        I1: Into<ExprExt>,
-        I2: Into<ExprExt>,
-    {
-        self.assert_zero_ext::<ExprExt, ExprExt>(x.into() - y.into());
-    }
-
-    fn assert_one_ext<ExprExt, I>(&mut self, x: I)
-    where
-        ExprExt: AbstractExtensionField<Self::Expr>,
-        I: Into<ExprExt>,
-    {
-        let xe: ExprExt = x.into();
-        let parts = xe.as_base_slice();
-        self.assert_one(parts[0].clone());
-        for part in &parts[1..] {
-            self.assert_zero(part.clone());
-        }
-    }
 }
 
 pub trait PairBuilder: AirBuilder {
     fn preprocessed(&self) -> Self::M;
 }
 
-pub trait PermutationAirBuilder: AirBuilder {
+pub trait ExtensionBuilder: AirBuilder {
     type EF: ExtensionField<Self::F>;
 
     type ExprEF: AbstractExtensionField<Self::Expr, F = Self::EF>;
 
     type VarEF: Into<Self::ExprEF> + Copy;
 
+    fn assert_zero_ext<I>(&mut self, x: I)
+    where
+        I: Into<Self::ExprEF>;
+
+    fn assert_eq_ext<I1, I2>(&mut self, x: I1, y: I2)
+    where
+        I1: Into<Self::ExprEF>,
+        I2: Into<Self::ExprEF>,
+    {
+        self.assert_zero_ext(x.into() - y.into());
+    }
+
+    fn assert_one_ext<I>(&mut self, x: I)
+    where
+        I: Into<Self::ExprEF>,
+    {
+        self.assert_eq_ext(x, Self::ExprEF::one())
+    }
+}
+
+pub trait PermutationAirBuilder: ExtensionBuilder {
     type MP: MatrixRowSlices<Self::VarEF>;
 
     fn permutation(&self) -> Self::MP;
@@ -165,21 +154,6 @@ pub struct FilteredAirBuilder<'a, AB: AirBuilder> {
     condition: AB::Expr,
 }
 
-impl<'a, AB: PermutationAirBuilder> PermutationAirBuilder for FilteredAirBuilder<'a, AB> {
-    type EF = AB::EF;
-    type VarEF = AB::VarEF;
-    type ExprEF = AB::ExprEF;
-    type MP = AB::MP;
-
-    fn permutation(&self) -> Self::MP {
-        self.inner.permutation()
-    }
-
-    fn permutation_randomness(&self) -> &[Self::EF] {
-        self.inner.permutation_randomness()
-    }
-}
-
 impl<'a, AB: AirBuilder> AirBuilder for FilteredAirBuilder<'a, AB> {
     type F = AB::F;
     type Expr = AB::Expr;
@@ -207,6 +181,32 @@ impl<'a, AB: AirBuilder> AirBuilder for FilteredAirBuilder<'a, AB> {
     }
 }
 
+impl<'a, AB: ExtensionBuilder> ExtensionBuilder for FilteredAirBuilder<'a, AB> {
+    type EF = AB::EF;
+    type VarEF = AB::VarEF;
+    type ExprEF = AB::ExprEF;
+
+    fn assert_zero_ext<I>(&mut self, x: I)
+    where
+        I: Into<Self::ExprEF>,
+    {
+        self.inner
+            .assert_zero_ext(x.into() * self.condition.clone());
+    }
+}
+
+impl<'a, AB: PermutationAirBuilder> PermutationAirBuilder for FilteredAirBuilder<'a, AB> {
+    type MP = AB::MP;
+
+    fn permutation(&self) -> Self::MP {
+        self.inner.permutation()
+    }
+
+    fn permutation_randomness(&self) -> &[Self::EF] {
+        self.inner.permutation_randomness()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use p3_matrix::MatrixRowSlices;

baby-bear/src/baby_bear.rs

@@ -263,10 +263,38 @@ impl TwoAdicField for BabyBear {
     const TWO_ADICITY: usize = 27;
 
     fn two_adic_generator(bits: usize) -> Self {
-        // TODO: Consider a `match` which may speed this up.
         assert!(bits <= Self::TWO_ADICITY);
-        let base = Self::from_canonical_u32(0x1a427a41); // generates the whole 2^TWO_ADICITY group
-        base.exp_power_of_2(Self::TWO_ADICITY - bits)
+        match bits {
+            0 => Self::one(),
+            1 => Self::from_canonical_u32(0x78000000),
+            2 => Self::from_canonical_u32(0x67055c21),
+            3 => Self::from_canonical_u32(0x5ee99486),
+            4 => Self::from_canonical_u32(0xbb4c4e4),
+            5 => Self::from_canonical_u32(0x2d4cc4da),
+            6 => Self::from_canonical_u32(0x669d6090),
+            7 => Self::from_canonical_u32(0x17b56c64),
+            8 => Self::from_canonical_u32(0x67456167),
+            9 => Self::from_canonical_u32(0x688442f9),
+            10 => Self::from_canonical_u32(0x145e952d),
+            11 => Self::from_canonical_u32(0x4fe61226),
+            12 => Self::from_canonical_u32(0x4c734715),
+            13 => Self::from_canonical_u32(0x11c33e2a),
+            14 => Self::from_canonical_u32(0x62c3d2b1),
+            15 => Self::from_canonical_u32(0x77cad399),
+            16 => Self::from_canonical_u32(0x54c131f4),
+            17 => Self::from_canonical_u32(0x4cabd6a6),
+            18 => Self::from_canonical_u32(0x5cf5713f),
+            19 => Self::from_canonical_u32(0x3e9430e8),
+            20 => Self::from_canonical_u32(0xba067a3),
+            21 => Self::from_canonical_u32(0x18adc27d),
+            22 => Self::from_canonical_u32(0x21fd55bc),
+            23 => Self::from_canonical_u32(0x4b859b3d),
+            24 => Self::from_canonical_u32(0x3bd57996),
+            25 => Self::from_canonical_u32(0x4483d85a),
+            26 => Self::from_canonical_u32(0x3a26eef8),
+            27 => Self::from_canonical_u32(0x1a427a41),
+            _ => unreachable!("Already asserted that bits <= Self::TWO_ADICITY"),
+        }
     }
 }
 
@@ -402,6 +430,17 @@ mod tests {
 
     type F = BabyBear;
 
+    #[test]
+    fn test_baby_bear_two_adicity_generators() {
+        let base = BabyBear::from_canonical_u32(0x1a427a41);
+        for bits in 0..=BabyBear::TWO_ADICITY {
+            assert_eq!(
+                BabyBear::two_adic_generator(bits),
+                base.exp_power_of_2(BabyBear::TWO_ADICITY - bits)
+            );
+        }
+    }
+
     #[test]
     fn test_baby_bear() {
         let f = F::from_canonical_u32(100);

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
+    }
 }

commit/src/adapters/extension_mmcs.rs

@@ -115,7 +115,9 @@ where
     InnerMat: Matrix<F>,
 {
     fn width(&self) -> usize {
-        self.inner.width() * <EF as AbstractExtensionField<F>>::D
+        let d = <EF as AbstractExtensionField<F>>::D;
+        debug_assert!(self.inner.width() % d == 0);
+        self.inner.width() / d
     }
 
     fn height(&self) -> usize {

commit/src/pcs.rs

@@ -2,6 +2,7 @@
 
 use alloc::vec;
 use alloc::vec::Vec;
+use core::fmt::Debug;
 
 use p3_challenger::FieldChallenger;
 use p3_field::{ExtensionField, Field};
@@ -25,7 +26,7 @@ pub trait Pcs<Val: Field, In: MatrixRows<Val>> {
     /// The opening argument.
     type Proof: Serialize + DeserializeOwned;
 
-    type Error;
+    type Error: Debug;
 
     fn commit_batches(&self, polynomials: Vec<In>) -> (Self::Commitment, Self::ProverData);
 

field/src/extension/binomial_extension.rs

@@ -13,13 +13,8 @@ use serde::{Deserialize, Serialize};
 use super::{HasFrobenius, HasTwoAdicBionmialExtension};
 use crate::extension::BinomiallyExtendable;
 use crate::field::Field;
-use crate::restriction::Res;
 use crate::{field_to_array, AbstractExtensionField, AbstractField, ExtensionField, TwoAdicField};
 
-/// The algebra F_D[X] / (X^D - W) where F_D is the binomial extension field over `F`.
-pub type BinomialExtensionAlgebra<F, const D: usize> =
-    BinomialExtensionField<Res<F, BinomialExtensionField<F, D>>, D>;
-
 #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, Serialize, Deserialize)]
 pub struct BinomialExtensionField<AF, const D: usize> {
     #[serde(

field/src/field.rs

@@ -242,6 +242,7 @@ pub trait PrimeField32: PrimeField64 {
 
 pub trait AbstractExtensionField<Base: AbstractField>:
     AbstractField
+    + From<Base>
     + Add<Base, Output = Self>
     + AddAssign<Base>
     + Sub<Base, Output = Self>