BabyBear with only Plonky3 changes

plonky3/Cargo.toml

@@ -17,15 +17,20 @@ p3-air = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-star
 p3-matrix = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 p3-field = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 p3-uni-stark = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
-p3-commit = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed", features = ["test-utils"] }
+p3-commit = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed", features = [
+  "test-utils",
+] }
 p3-poseidon2 = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 p3-poseidon = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 p3-fri = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 # We don't use p3-maybe-rayon directly, but it is a dependency of p3-uni-stark.
 # Activating the "parallel" feature gives us parallelism in the prover.
-p3-maybe-rayon = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed", features = ["parallel"] }
+p3-maybe-rayon = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed", features = [
+  "parallel",
+] }
 p3-mds = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 p3-merkle-tree = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
+p3-baby-bear = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 p3-goldilocks = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 p3-symmetric = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }
 p3-dft = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "uni-stark-with-fixed" }

plonky3/src/baby_bear.rs

@@ -0,0 +1,121 @@
+//! The concrete parameters used in the prover
+//! Inspired from [this example](https://github.com/Plonky3/Plonky3/blob/6a1b0710fdf85136d0fdd645b92933615867740a/keccak-air/examples/prove_goldilocks_keccak.rs#L57)
+
+use lazy_static::lazy_static;
+
+use crate::params::FieldElementMap;
+use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear, MdsMatrixBabyBear};
+use p3_challenger::DuplexChallenger;
+use p3_commit::ExtensionMmcs;
+use p3_dft::Radix2DitParallel;
+use p3_field::{extension::BinomialExtensionField, Field};
+use p3_fri::{FriConfig, TwoAdicFriPcs};
+use p3_merkle_tree::FieldMerkleTreeMmcs;
+use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
+use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
+use p3_uni_stark::StarkConfig;
+
+use rand::{distributions::Standard, Rng, SeedableRng};
+
+use powdr_number::{BabyBearField, FieldElement};
+
+const DEGREE: usize = 2;
+type Challenge<T: FieldElementMap> = BinomialExtensionField<T::P3Field, DEGREE>;
+
+const WIDTH: usize = 16;
+const D: u64 = 7;
+type Perm = Poseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, WIDTH, D>;
+
+// const ROUNDS_F: usize = poseidon2_round_numbers_128::<dyn BabyBear>(WIDTH, D).0;
+// const ROUNDS_P: usize = poseidon2_round_numbers_128::<dyn BabyBear>(WIDTH, D).1;
+
+// params directly taken from plonky3's poseidon2_round_numbers_128 function
+// to guarentee 128-bit security.
+
+const ROUNDS_F: usize = 8;
+const ROUNDS_P: usize = 13;
+
+const RATE: usize = 8;
+const OUT: usize = 8;
+type Hash<T: FieldElementMap> = PaddingFreeSponge<Perm, WIDTH, RATE, OUT>;
+
+const N: usize = 2;
+const CHUNK: usize = 8;
+type Compress<T: FieldElementMap> = TruncatedPermutation<Perm, N, CHUNK, WIDTH>;
+
+const DIGEST_ELEMS: usize = 8;
+type ValMmcs<T: FieldElementMap> = FieldMerkleTreeMmcs<
+    <T::P3Field as Field>::Packing,
+    <T::P3Field as Field>::Packing,
+    Hash<T>,
+    Compress<T>,
+    DIGEST_ELEMS,
+>;
+
+pub type Val = BabyBear;
+
+pub type Challenger<T: FieldElementMap> = DuplexChallenger<T::P3Field, Perm, WIDTH, RATE>;
+type ChallengeMmcs<T: FieldElementMap> = ExtensionMmcs<T::P3Field, Challenge<T>, ValMmcs<T>>;
+type Dft = Radix2DitParallel;
+type MyPcs<T: FieldElementMap> = TwoAdicFriPcs<T::P3Field, Dft, ValMmcs<T>, ChallengeMmcs<T>>;
+pub type Config<T: FieldElementMap> = StarkConfig<MyPcs<T>, Challenge<T>, Challenger<T>>;
+
+const FRI_LOG_BLOWUP: usize = 1;
+const FRI_NUM_QUERIES: usize = 100;
+const FRI_PROOF_OF_WORK_BITS: usize = 16;
+
+const RNG_SEED: u64 = 42;
+
+lazy_static! {
+    static ref PERM_BB: Perm = Perm::new(
+        ROUNDS_F,
+        rand_chacha::ChaCha8Rng::seed_from_u64(RNG_SEED)
+            .sample_iter(Standard)
+            .take(ROUNDS_F)
+            .collect::<Vec<[BabyBear; WIDTH]>>(),
+        Poseidon2ExternalMatrixGeneral,
+        ROUNDS_P,
+        rand_chacha::ChaCha8Rng::seed_from_u64(RNG_SEED)
+            .sample_iter(Standard)
+            .take(ROUNDS_P)
+            .collect(),
+        DiffusionMatrixBabyBear::default()
+    );
+}
+
+impl FieldElementMap for BabyBearField {
+    type P3Field = BabyBear;
+    type MdsMatrix = MdsMatrixBabyBear;
+
+    fn to_p3_field(&self) -> Self::P3Field {
+        BabyBear::from_canonical_u32(self.to_integer().try_into_u32().unwrap())
+    }
+}
+
+pub(crate) fn get_challenger_baby_bear<T: FieldElementMap>() -> Challenger<T> {
+    Challenger::new(PERM_BB.clone())
+}
+
+pub(crate) fn get_config_baby_bear<T: FieldElementMap>(
+) -> StarkConfig<MyPcs<T>, Challenge<T>, Challenger<T>> {
+    let hash = Hash::<T>::new(PERM_BB.clone());
+
+    let compress = Compress::<T>::new(PERM_BB.clone());
+
+    let val_mmcs = ValMmcs::<T>::new(hash, compress);
+
+    let challenge_mmcs = ChallengeMmcs::<T>::new(val_mmcs.clone());
+
+    let dft = Dft {};
+
+    let fri_config = FriConfig {
+        log_blowup: FRI_LOG_BLOWUP,
+        num_queries: FRI_NUM_QUERIES,
+        proof_of_work_bits: FRI_PROOF_OF_WORK_BITS,
+        mmcs: challenge_mmcs,
+    };
+
+    let pcs = MyPcs::<T>::new(dft, val_mmcs, fri_config);
+
+    Config::new(pcs)
+}

plonky3/src/circuit_builder.rs

@@ -6,22 +6,24 @@
 //! everywhere save for at row j is constructed to constrain s * (pub - x) on
 //! every row.
 
-use std::{any::TypeId, collections::BTreeMap};
+use std::collections::BTreeMap;
 
+use crate::params::FieldElementMap;
 use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder};
-use p3_field::AbstractField;
-use p3_goldilocks::Goldilocks;
+// use p3_baby_bear::BabyBear;
+// use p3_field::{extension::ComplexExtendable, AbstractField, PrimeField};
+// use p3_goldilocks::Goldilocks;
 use p3_matrix::{dense::RowMajorMatrix, Matrix};
+// use p3_mds::MdsPermutation;
 use powdr_ast::analyzed::{
     AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression,
     AlgebraicUnaryOperation, AlgebraicUnaryOperator, Analyzed, IdentityKind, PolynomialType,
 };
 use powdr_executor::witgen::WitgenCallback;
-use powdr_number::{FieldElement, GoldilocksField, LargeInt};
+// use powdr_number::{BabyBearField, FieldElement, GoldilocksField, LargeInt};
+use powdr_number::FieldElement;
 
-pub type Val = p3_goldilocks::Goldilocks;
-
-pub(crate) struct PowdrCircuit<'a, T> {
+pub(crate) struct PowdrCircuit<'a, T: FieldElementMap> {
     /// The analyzed PIL
     analyzed: &'a Analyzed<T>,
     /// The value of the witness columns, if set
@@ -30,11 +32,11 @@ pub(crate) struct PowdrCircuit<'a, T> {
     _witgen_callback: Option<WitgenCallback<T>>,
     /// The matrix of preprocessed values, used in debug mode to check the constraints before proving
     #[cfg(debug_assertions)]
-    preprocessed: Option<RowMajorMatrix<Goldilocks>>,
+    preprocessed: Option<RowMajorMatrix<<T as FieldElementMap>::P3Field>>,
 }
 
-impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
-    pub fn generate_trace_rows(&self) -> RowMajorMatrix<Goldilocks> {
+impl<'a, T: FieldElement + FieldElementMap + Sync> PowdrCircuit<'a, T> {
+    pub fn generate_trace_rows(&self) -> RowMajorMatrix<<T as FieldElementMap>::P3Field> {
         // an iterator over all columns, committed then fixed
         let witness = self.witness().iter();
         let degrees = self.analyzed.degrees();
@@ -48,7 +50,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
                         // witness values
                         witness.clone().map(move |(_, v)| v[i as usize])
                     })
-                    .map(cast_to_goldilocks)
+                    .map(|f| f.to_p3_field())
                     .collect()
             }
             0 => {
@@ -62,12 +64,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
     }
 }
 
-pub fn cast_to_goldilocks<T: FieldElement>(v: T) -> Val {
-    assert_eq!(TypeId::of::<T>(), TypeId::of::<GoldilocksField>());
-    Val::from_canonical_u64(v.to_integer().try_into_u64().unwrap())
-}
-
-impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
+impl<'a, T: FieldElement + FieldElementMap> PowdrCircuit<'a, T> {
     pub(crate) fn new(analyzed: &'a Analyzed<T>) -> Self {
         if analyzed
             .definitions
@@ -91,7 +88,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
     }
 
     /// Calculates public values from generated witness values.
-    pub(crate) fn get_public_values(&self) -> Vec<Goldilocks> {
+    pub(crate) fn get_public_values(&self) -> Vec<<T as FieldElementMap>::P3Field> {
         let witness = self
             .witness
             .as_ref()
@@ -105,7 +102,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
             .iter()
             .map(|(col_name, _, idx)| {
                 let vals = *witness.get(&col_name).unwrap();
-                cast_to_goldilocks(vals[*idx])
+                vals[*idx].to_p3_field()
             })
             .collect()
     }
@@ -128,14 +125,16 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
     #[cfg(debug_assertions)]
     pub(crate) fn with_preprocessed(
         mut self,
-        preprocessed_matrix: RowMajorMatrix<Goldilocks>,
+        preprocessed_matrix: RowMajorMatrix<<T as FieldElementMap>::P3Field>,
     ) -> Self {
         self.preprocessed = Some(preprocessed_matrix);
         self
     }
 
     /// Conversion to plonky3 expression
-    fn to_plonky3_expr<AB: AirBuilder<F = Val> + AirBuilderWithPublicValues>(
+    fn to_plonky3_expr<
+        AB: AirBuilder<F = <T as FieldElementMap>::P3Field> + AirBuilderWithPublicValues,
+    >(
         &self,
         e: &AlgebraicExpression<T>,
         main: &AB::M,
@@ -172,7 +171,7 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
                 .get(id)
                 .expect("Referenced public value does not exist"))
             .into(),
-            AlgebraicExpression::Number(n) => AB::Expr::from(cast_to_goldilocks(*n)),
+            AlgebraicExpression::Number(n) => AB::Expr::from(n.to_p3_field()),
             AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) => {
                 let left = self.to_plonky3_expr::<AB>(left, main, fixed, publics);
                 let right = self.to_plonky3_expr::<AB>(right, main, fixed, publics);
@@ -204,7 +203,9 @@ impl<'a, T: FieldElement> PowdrCircuit<'a, T> {
 
 /// An extension of [Air] allowing access to the number of fixed columns
 
-impl<'a, T: FieldElement> BaseAir<Val> for PowdrCircuit<'a, T> {
+impl<'a, T: FieldElement + FieldElementMap> BaseAir<<T as FieldElementMap>::P3Field>
+    for PowdrCircuit<'a, T>
+{
     fn width(&self) -> usize {
         self.analyzed.commitment_count()
     }
@@ -213,7 +214,7 @@ impl<'a, T: FieldElement> BaseAir<Val> for PowdrCircuit<'a, T> {
         self.analyzed.constant_count() + self.analyzed.publics_count()
     }
 
-    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<Val>> {
+    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<<T as FieldElementMap>::P3Field>> {
         #[cfg(debug_assertions)]
         {
             self.preprocessed.clone()
@@ -223,8 +224,11 @@ impl<'a, T: FieldElement> BaseAir<Val> for PowdrCircuit<'a, T> {
     }
 }
 
-impl<'a, T: FieldElement, AB: AirBuilderWithPublicValues<F = Val> + PairBuilder> Air<AB>
-    for PowdrCircuit<'a, T>
+impl<
+        'a,
+        T: FieldElement + FieldElementMap,
+        AB: AirBuilderWithPublicValues<F = <T as FieldElementMap>::P3Field> + PairBuilder,
+    > Air<AB> for PowdrCircuit<'a, T>
 {
     fn eval(&self, builder: &mut AB) {
         let main = builder.main();