BabyBear with only Plonky3 changes (powdr-labs#1693)

Adding BabyBear params to Plonky3 crate.

---------

Co-authored-by: Thibaut Schaeffer <[email protected]>
Co-authored-by: Lucas Clemente Vella <[email protected]>
Co-authored-by: schaeff <[email protected]>
Co-authored-by: Leo Alt <[email protected]>

backend/src/plonky3/mod.rs

@@ -1,32 +1,70 @@
-use std::{io, path::PathBuf, sync::Arc};
+use std::{
+    any::{Any, TypeId},
+    io,
+    path::PathBuf,
+    sync::Arc,
+};
 
 use powdr_ast::analyzed::Analyzed;
 use powdr_executor::{
     constant_evaluator::{get_uniquely_sized_cloned, VariablySizedColumn},
     witgen::WitgenCallback,
 };
-use powdr_number::{FieldElement, GoldilocksField, LargeInt};
-use powdr_plonky3::Plonky3Prover;
+use powdr_number::{BabyBearField, FieldElement, GoldilocksField};
+use powdr_plonky3::{Commitment, FieldElementMap, Plonky3Prover, ProverData};
 
 use crate::{Backend, BackendFactory, BackendOptions, Error, Proof};
 
 pub(crate) struct Factory;
 
+fn try_create<FInner: FieldElementMap, FOuter: FieldElement>(
+    pil: &Arc<Analyzed<FOuter>>,
+    fixed: &Arc<Vec<(String, Vec<FOuter>)>>,
+    verification_key: &mut Option<&mut dyn io::Read>,
+) -> Option<Box<dyn Backend<FOuter>>>
+where
+    ProverData<FInner>: Send,
+    Commitment<FInner>: Send,
+{
+    // We ensure that FInner and FOuter are the same type, so we can even safely
+    // transmute between them.
+    if TypeId::of::<FInner>() != TypeId::of::<FOuter>() {
+        return None;
+    }
+
+    let pil = (pil as &dyn Any)
+        .downcast_ref::<Arc<Analyzed<FInner>>>()
+        .unwrap();
+    let fixed = (fixed as &dyn Any)
+        .downcast_ref::<Arc<Vec<(String, Vec<FInner>)>>>()
+        .unwrap();
+
+    let mut p3 = Box::new(Plonky3Prover::new(pil.clone(), fixed.clone()));
+
+    if let Some(verification_key) = verification_key {
+        p3.set_verifying_key(*verification_key);
+    } else {
+        p3.setup();
+    }
+
+    let p3: Box<dyn Backend<FInner>> = p3;
+    let p3 = Box::into_raw(p3);
+
+    // This is safe because we know that FInner == FOuter.
+    Some(unsafe { Box::from_raw(p3 as *mut dyn Backend<FOuter>) })
+}
+
 impl<T: FieldElement> BackendFactory<T> for Factory {
     fn create(
         &self,
         pil: Arc<Analyzed<T>>,
         fixed: Arc<Vec<(String, VariablySizedColumn<T>)>>,
         _output_dir: Option<PathBuf>,
         setup: Option<&mut dyn io::Read>,
-        verification_key: Option<&mut dyn io::Read>,
+        mut verification_key: Option<&mut dyn io::Read>,
         verification_app_key: Option<&mut dyn io::Read>,
         _: BackendOptions,
     ) -> Result<Box<dyn crate::Backend<T>>, Error> {
-        if T::modulus().to_arbitrary_integer() != GoldilocksField::modulus().to_arbitrary_integer()
-        {
-            unimplemented!("plonky3 is only implemented for the Goldilocks field");
-        }
         if setup.is_some() {
             return Err(Error::NoSetupAvailable);
         }
@@ -41,19 +79,26 @@ impl<T: FieldElement> BackendFactory<T> for Factory {
             get_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
         );
 
-        let mut p3 = Box::new(Plonky3Prover::new(pil, fixed));
-
-        if let Some(verification_key) = verification_key {
-            p3.set_verifying_key(verification_key);
-        } else {
-            p3.setup();
-        }
-
-        Ok(p3)
+        Ok(
+            if let Some(p3) = try_create::<GoldilocksField, T>(&pil, &fixed, &mut verification_key)
+            {
+                p3
+            } else if let Some(p3) =
+                try_create::<BabyBearField, T>(&pil, &fixed, &mut verification_key)
+            {
+                p3
+            } else {
+                unimplemented!("unsupported field type: {:?}", TypeId::of::<T>())
+            },
+        )
     }
 }
 
-impl<T: FieldElement> Backend<T> for Plonky3Prover<T> {
+impl<T: FieldElementMap> Backend<T> for Plonky3Prover<T>
+where
+    ProverData<T>: Send,
+    Commitment<T>: Send,
+{
     fn verify(&self, proof: &[u8], instances: &[Vec<T>]) -> Result<(), Error> {
         Ok(self.verify(proof, instances)?)
     }

number/Cargo.toml

@@ -13,8 +13,8 @@ ark-bn254 = { version = "0.4.0", default-features = false, features = [
 ] }
 ark-ff = "0.4.2"
 ark-serialize = "0.4.2"
-p3-baby-bear = { 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-baby-bear = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
+p3-field = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 num-bigint = { version = "0.4.3", features = ["serde"] }
 num-traits = "0.2.15"
 csv = "1.3"

number/src/baby_bear.rs

@@ -48,6 +48,10 @@ impl BabyBearField {
     fn to_canonical_u32(self) -> u32 {
         self.0.as_canonical_u32()
     }
+
+    pub fn into_inner(self) -> BabyBear {
+        self.0
+    }
 }
 
 impl FieldElement for BabyBearField {

pipeline/src/test_util.rs

@@ -77,12 +77,6 @@ pub fn regular_test(file_name: &str, inputs: &[i32]) {
     pipeline_bb.compute_witness().unwrap();
 }
 
-pub fn regular_test_only_babybear(file_name: &str, inputs: &[i32]) {
-    let inputs_bb = inputs.iter().map(|x| BabyBearField::from(*x)).collect();
-    let mut pipeline_bb = make_prepared_pipeline(file_name, inputs_bb, vec![]);
-    pipeline_bb.compute_witness().unwrap();
-}
-
 pub fn regular_test_without_babybear(file_name: &str, inputs: &[i32]) {
     let inputs_gl = inputs.iter().map(|x| GoldilocksField::from(*x)).collect();
     let pipeline_gl = make_prepared_pipeline(file_name, inputs_gl, vec![]);
@@ -296,9 +290,9 @@ pub fn gen_halo2_proof(pipeline: Pipeline<Bn254Field>, backend: BackendVariant)
 pub fn gen_halo2_proof(_pipeline: Pipeline<Bn254Field>, _backend: BackendVariant) {}
 
 #[cfg(feature = "plonky3")]
-pub fn test_plonky3_with_backend_variant(
+pub fn test_plonky3_with_backend_variant<T: FieldElement>(
     file_name: &str,
-    inputs: Vec<GoldilocksField>,
+    inputs: Vec<T>,
     backend: BackendVariant,
 ) {
     let backend = match backend {
@@ -314,7 +308,7 @@ pub fn test_plonky3_with_backend_variant(
     // Generate a proof
     let proof = pipeline.compute_proof().cloned().unwrap();
 
-    let publics: Vec<GoldilocksField> = pipeline
+    let publics: Vec<T> = pipeline
         .publics()
         .clone()
         .unwrap()
@@ -341,10 +335,10 @@ pub fn test_plonky3_with_backend_variant(
 }
 
 #[cfg(not(feature = "plonky3"))]
-pub fn test_plonky3_with_backend_variant(_: &str, _: Vec<GoldilocksField>, _: BackendVariant) {}
+pub fn test_plonky3_with_backend_variant<T: FieldElement>(_: &str, _: Vec<T>, _: BackendVariant) {}
 
 #[cfg(not(feature = "plonky3"))]
-pub fn gen_plonky3_proof(_: &str, _: Vec<GoldilocksField>) {}
+pub fn gen_plonky3_proof<T: FieldElement>(_: &str, _: Vec<T>) {}
 
 /// Returns the analyzed PIL containing only the std library.
 pub fn std_analyzed<T: FieldElement>() -> Analyzed<T> {

pipeline/tests/asm.rs

@@ -39,7 +39,11 @@ fn simple_sum_asm() {
     let f = "asm/simple_sum.asm";
     let i = [16, 4, 1, 2, 8, 5];
     regular_test(f, &i);
-    test_plonky3_with_backend_variant(f, slice_to_vec(&i), BackendVariant::Composite);
+    test_plonky3_with_backend_variant::<GoldilocksField>(
+        f,
+        slice_to_vec(&i),
+        BackendVariant::Composite,
+    );
 }
 
 #[test]

pipeline/tests/pil.rs

@@ -87,7 +87,11 @@ fn permutation_with_selector() {
 fn fibonacci() {
     let f = "pil/fibonacci.pil";
     regular_test(f, Default::default());
-    test_plonky3_with_backend_variant(f, Default::default(), BackendVariant::Monolithic);
+    test_plonky3_with_backend_variant::<GoldilocksField>(
+        f,
+        Default::default(),
+        BackendVariant::Monolithic,
+    );
 }
 
 #[test]
@@ -241,7 +245,11 @@ fn halo_without_lookup() {
 #[test]
 fn add() {
     let f = "pil/add.pil";
-    test_plonky3_with_backend_variant(f, Default::default(), BackendVariant::Monolithic);
+    test_plonky3_with_backend_variant::<GoldilocksField>(
+        f,
+        Default::default(),
+        BackendVariant::Monolithic,
+    );
 }
 
 #[test]

pipeline/tests/powdr_std.rs

@@ -6,8 +6,8 @@ use powdr_pil_analyzer::evaluator::Value;
 use powdr_pipeline::{
     test_util::{
         evaluate_function, evaluate_integer_function, execute_test_file, gen_estark_proof,
-        gen_halo2_proof, make_simple_prepared_pipeline, regular_test, regular_test_only_babybear,
-        std_analyzed, test_halo2, test_pilcom, BackendVariant,
+        gen_halo2_proof, make_simple_prepared_pipeline, regular_test, std_analyzed, test_halo2,
+        test_pilcom, test_plonky3_with_backend_variant, BackendVariant,
     },
     Pipeline,
 };
@@ -186,14 +186,14 @@ fn binary_test() {
 #[ignore = "Too slow"]
 fn binary_bb_8_test() {
     let f = "std/binary_bb_test_8.asm";
-    regular_test_only_babybear(f, &[]);
+    test_plonky3_with_backend_variant::<BabyBearField>(f, vec![], BackendVariant::Composite);
 }
 
 #[test]
 #[ignore = "Too slow"]
 fn binary_bb_16_test() {
     let f = "std/binary_bb_test_16.asm";
-    regular_test_only_babybear(f, &[]);
+    test_plonky3_with_backend_variant::<BabyBearField>(f, vec![], BackendVariant::Composite);
 }
 
 #[test]

plonky3/Cargo.toml

@@ -18,18 +18,19 @@ p3-matrix = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main"
 p3-field = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 p3-uni-stark = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 p3-commit = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main", features = [
-    "test-utils",
+  "test-utils",
 ] }
 p3-poseidon2 = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 p3-poseidon = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 p3-fri = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 # 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 = "main", features = [
-    "parallel",
+  "parallel",
 ] }
 p3-mds = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 p3-merkle-tree = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
+p3-baby-bear = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 p3-goldilocks = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 p3-symmetric = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }
 p3-dft = { git = "https://github.com/powdr-labs/Plonky3.git", branch = "main" }

plonky3/src/baby_bear.rs

@@ -0,0 +1,109 @@
+//! The concrete parameters used in the prover
+//! Inspired from [this example](https://github.com/Plonky3/Plonky3/blob/6a1b0710fdf85136d0fdd645b92933615867740a/keccak-air/examples/prove_baby_bear_poseidon2.rs)
+
+use lazy_static::lazy_static;
+
+use crate::params::{Challenger, FieldElementMap, Plonky3Field};
+use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
+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;
+
+const D: u64 = 7;
+// 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 WIDTH: usize = 16;
+type Perm = Poseidon2<BabyBear, Poseidon2ExternalMatrixGeneral, DiffusionMatrixBabyBear, WIDTH, D>;
+
+const DEGREE: usize = 4;
+type FriChallenge = BinomialExtensionField<BabyBear, DEGREE>;
+
+const RATE: usize = 8;
+const OUT: usize = 8;
+type FriChallenger = DuplexChallenger<BabyBear, Perm, WIDTH, RATE>;
+type Hash = PaddingFreeSponge<Perm, WIDTH, RATE, OUT>;
+
+const N: usize = 2;
+const CHUNK: usize = 8;
+type Compress = TruncatedPermutation<Perm, N, CHUNK, WIDTH>;
+const DIGEST_ELEMS: usize = 8;
+type ValMmcs = FieldMerkleTreeMmcs<
+    <BabyBear as Field>::Packing,
+    <BabyBear as Field>::Packing,
+    Hash,
+    Compress,
+    DIGEST_ELEMS,
+>;
+
+type ChallengeMmcs = ExtensionMmcs<BabyBear, FriChallenge, ValMmcs>;
+type Dft = Radix2DitParallel;
+type MyPcs = TwoAdicFriPcs<BabyBear, Dft, ValMmcs, ChallengeMmcs>;
+
+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 Config = StarkConfig<MyPcs, FriChallenge, FriChallenger>;
+    fn into_p3_field(self) -> Plonky3Field<Self> {
+        self.into_inner()
+    }
+
+    fn get_challenger() -> Challenger<Self> {
+        FriChallenger::new(PERM_BB.clone())
+    }
+
+    fn get_config() -> Self::Config {
+        let hash = Hash::new(PERM_BB.clone());
+
+        let compress = Compress::new(PERM_BB.clone());
+
+        let val_mmcs = ValMmcs::new(hash, compress);
+
+        let challenge_mmcs = ChallengeMmcs::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::new(dft, val_mmcs, fri_config);
+
+        Self::Config::new(pcs)
+    }
+}