Plonky3 + Goldilocks: Use Poseidon2 instead of Poseidon

powdr-labs · Oct 9, 2024 · af7f404 · af7f404
1 parent 5d1a2fa
commit af7f404
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 24 deletions.
diff --git a/plonky3/src/params/baby_bear.rs b/plonky3/src/params/baby_bear.rs
@@ -11,7 +11,7 @@ use p3_dft::Radix2DitParallel;
 use p3_field::{extension::BinomialExtensionField, Field};
 use p3_fri::{FriConfig, TwoAdicFriPcs};
 use p3_merkle_tree::MerkleTreeMmcs;
-use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
+use p3_poseidon2::{poseidon2_round_numbers_128, Poseidon2, Poseidon2ExternalMatrixGeneral};
 use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
 use p3_uni_stark::StarkConfig;
 
@@ -20,10 +20,6 @@ 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 guarantee 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>;
 
@@ -58,17 +54,20 @@ const FRI_PROOF_OF_WORK_BITS: usize = 16;
 const RNG_SEED: u64 = 42;
 
 lazy_static! {
+    static ref ROUNDS: (usize, usize) = poseidon2_round_numbers_128::<Goldilocks>(WIDTH, D);
+    static ref ROUNDS_F: usize = ROUNDS.0;
+    static ref ROUNDS_P: usize = ROUNDS.1;
     static ref PERM_BB: Perm = Perm::new(
-        ROUNDS_F,
+        *ROUNDS_F,
         rand_chacha::ChaCha8Rng::seed_from_u64(RNG_SEED)
             .sample_iter(Standard)
-            .take(ROUNDS_F)
+            .take(*ROUNDS_F)
             .collect::<Vec<[BabyBear; WIDTH]>>(),
         Poseidon2ExternalMatrixGeneral,
-        ROUNDS_P,
+        *ROUNDS_P,
         rand_chacha::ChaCha8Rng::seed_from_u64(RNG_SEED)
             .sample_iter(Standard)
-            .take(ROUNDS_P)
+            .take(*ROUNDS_P)
             .collect(),
         DiffusionMatrixBabyBear::default()
     );

diff --git a/plonky3/src/params/goldilocks.rs b/plonky3/src/params/goldilocks.rs
@@ -1,27 +1,30 @@
 //! 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)
+//! (But using Poseidon2 instead of Poseidon)
 
 use lazy_static::lazy_static;
+use p3_poseidon2::{poseidon2_round_numbers_128, Poseidon2, Poseidon2ExternalMatrixGeneral};
 
 use crate::params::{Challenger, FieldElementMap, Plonky3Field};
 use p3_challenger::DuplexChallenger;
 use p3_commit::ExtensionMmcs;
 use p3_dft::Radix2DitParallel;
 use p3_field::{extension::BinomialExtensionField, AbstractField, Field, PrimeField64};
 use p3_fri::{FriConfig, TwoAdicFriPcs};
-use p3_goldilocks::{Goldilocks, MdsMatrixGoldilocks};
+use p3_goldilocks::{DiffusionMatrixGoldilocks, Goldilocks};
 use p3_merkle_tree::MerkleTreeMmcs;
-use p3_poseidon::Poseidon;
 use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
 use p3_uni_stark::StarkConfig;
 use powdr_number::{FieldElement, GoldilocksField, LargeInt};
 use rand::{distributions::Standard, Rng, SeedableRng};
 
+const D: u64 = 7;
+const WIDTH: usize = 16;
+type Perm =
+    Poseidon2<Goldilocks, Poseidon2ExternalMatrixGeneral, DiffusionMatrixGoldilocks, WIDTH, D>;
+
 const DEGREE: usize = 2;
 type FriChallenge = BinomialExtensionField<Goldilocks, DEGREE>;
-const WIDTH: usize = 8;
-const ALPHA: u64 = 7;
-type Perm = Poseidon<Goldilocks, MdsMatrixGoldilocks, WIDTH, ALPHA>;
 
 const RATE: usize = 4;
 const OUT: usize = 4;
@@ -45,27 +48,29 @@ type ChallengeMmcs = ExtensionMmcs<Goldilocks, FriChallenge, ValMmcs>;
 type Dft = Radix2DitParallel<Goldilocks>;
 type MyPcs = TwoAdicFriPcs<Goldilocks, Dft, ValMmcs, ChallengeMmcs>;
 
-const HALF_NUM_FULL_ROUNDS: usize = 4;
-const NUM_PARTIAL_ROUNDS: usize = 22;
-
 const FRI_LOG_BLOWUP: usize = 1;
 const FRI_NUM_QUERIES: usize = 100;
 const FRI_PROOF_OF_WORK_BITS: usize = 16;
 
-const NUM_ROUNDS: usize = 2 * HALF_NUM_FULL_ROUNDS + NUM_PARTIAL_ROUNDS;
-const NUM_CONSTANTS: usize = WIDTH * NUM_ROUNDS;
-
 const RNG_SEED: u64 = 42;
 
 lazy_static! {
+    static ref ROUNDS: (usize, usize) = poseidon2_round_numbers_128::<Goldilocks>(WIDTH, D);
+    static ref ROUNDS_F: usize = ROUNDS.0;
+    static ref ROUNDS_P: usize = ROUNDS.1;
     static ref PERM_GL: Perm = Perm::new(
-        HALF_NUM_FULL_ROUNDS,
-        NUM_PARTIAL_ROUNDS,
+        *ROUNDS_F,
+        rand_chacha::ChaCha8Rng::seed_from_u64(RNG_SEED)
+            .sample_iter(Standard)
+            .take(*ROUNDS_F)
+            .collect::<Vec<[Goldilocks; WIDTH]>>(),
+        Poseidon2ExternalMatrixGeneral,
+        *ROUNDS_P,
         rand_chacha::ChaCha8Rng::seed_from_u64(RNG_SEED)
             .sample_iter(Standard)
-            .take(NUM_CONSTANTS)
+            .take(*ROUNDS_P)
             .collect(),
-        MdsMatrixGoldilocks,
+        DiffusionMatrixGoldilocks::default()
     );
 }