feat: add new permutation-related APIs

crates/ff_ext/src/babybear.rs

@@ -1,5 +1,5 @@
 pub mod impl_babybear {
-    use crate::array_try_from_uniform_bytes;
+    use crate::{array_try_from_uniform_bytes, wrapper::Wrapper};
     use p3::{
         self,
         babybear::{BabyBear, Poseidon2BabyBear},
@@ -94,20 +94,14 @@ pub mod impl_babybear {
 
     #[cfg(debug_assertions)]
     use crate::poseidon::impl_instruments::*;
-    #[cfg(debug_assertions)]
-    use p3::symmetric::CryptographicPermutation;
 
-    #[cfg(debug_assertions)]
-    impl CryptographicPermutation<[BabyBear; POSEIDON2_BABYBEAR_WIDTH]>
-        for Instrumented<Poseidon2BabyBear<POSEIDON2_BABYBEAR_WIDTH>>
-    {
-    }
+    type WP = Wrapper<Poseidon2BabyBear<POSEIDON2_BABYBEAR_WIDTH>, POSEIDON2_BABYBEAR_WIDTH>;
 
     impl PoseidonField for BabyBear {
         #[cfg(debug_assertions)]
-        type P = Instrumented<Poseidon2BabyBear<POSEIDON2_BABYBEAR_WIDTH>>;
+        type P = Instrumented<WP>;
         #[cfg(not(debug_assertions))]
-        type P = Poseidon2BabyBear<POSEIDON2_BABYBEAR_WIDTH>;
+        type P = WP;
 
         type T = DuplexChallenger<Self, Self::P, POSEIDON2_BABYBEAR_WIDTH, POSEIDON2_BABYBEAR_RATE>;
         type S = PaddingFreeSponge<Self::P, POSEIDON2_BABYBEAR_WIDTH, POSEIDON2_BABYBEAR_RATE, 8>;
@@ -124,24 +118,34 @@ pub mod impl_babybear {
 
         #[cfg(debug_assertions)]
         fn get_default_perm() -> Self::P {
-            Instrumented::new(Poseidon2BabyBear::new(
+            Instrumented::new(Wrapper::new(Poseidon2BabyBear::new(
                 ExternalLayerConstants::new(
                     BABYBEAR_RC16_EXTERNAL_INITIAL.to_vec(),
                     BABYBEAR_RC16_EXTERNAL_FINAL.to_vec(),
                 ),
                 BABYBEAR_RC16_INTERNAL.to_vec(),
-            ))
+            )))
         }
 
         #[cfg(not(debug_assertions))]
         fn get_default_perm() -> Self::P {
-            Poseidon2BabyBear::new(
+            Wrapper::new(Poseidon2BabyBear::new(
                 ExternalLayerConstants::new(
                     BABYBEAR_RC16_EXTERNAL_INITIAL.to_vec(),
                     BABYBEAR_RC16_EXTERNAL_FINAL.to_vec(),
                 ),
                 BABYBEAR_RC16_INTERNAL.to_vec(),
-            )
+            ))
+        }
+
+        fn get_default_perm_rc() -> Vec<Self> {
+            BABYBEAR_RC16_EXTERNAL_INITIAL
+                .iter()
+                .flatten()
+                .chain(BABYBEAR_RC16_INTERNAL.iter())
+                .chain(BABYBEAR_RC16_EXTERNAL_FINAL.iter().flatten())
+                .cloned()
+                .collect()
         }
 
         fn get_default_sponge() -> Self::S {

crates/ff_ext/src/goldilock.rs

@@ -3,6 +3,7 @@ pub mod impl_goldilocks {
         ExtensionField, FieldFrom, FieldInto, FromUniformBytes, SmallField,
         array_try_from_uniform_bytes, impl_from_uniform_bytes_for_binomial_extension,
         poseidon::{PoseidonField, new_array},
+        wrapper::Wrapper,
     };
     use p3::{
         challenger::DuplexChallenger,
@@ -21,8 +22,6 @@ pub mod impl_goldilocks {
 
     #[cfg(debug_assertions)]
     use crate::poseidon::impl_instruments::*;
-    #[cfg(debug_assertions)]
-    use p3::symmetric::CryptographicPermutation;
 
     pub type GoldilocksExt2 = BinomialExtensionField<Goldilocks, 2>;
 
@@ -47,17 +46,13 @@ pub mod impl_goldilocks {
     pub const POSEIDON2_GOLDILICK_WIDTH: usize = 8;
     pub const POSEIDON2_GOLDILICK_RATE: usize = 4;
 
-    #[cfg(debug_assertions)]
-    impl CryptographicPermutation<[Goldilocks; POSEIDON2_GOLDILICK_WIDTH]>
-        for Instrumented<Poseidon2GoldilocksHL<POSEIDON2_GOLDILICK_WIDTH>>
-    {
-    }
-
+    type WP = Wrapper<Poseidon2GoldilocksHL<POSEIDON2_GOLDILICK_WIDTH>, POSEIDON2_GOLDILICK_WIDTH>;
     impl PoseidonField for Goldilocks {
         #[cfg(debug_assertions)]
-        type P = Instrumented<Poseidon2GoldilocksHL<POSEIDON2_GOLDILICK_WIDTH>>;
+        type P = Instrumented<WP>;
         #[cfg(not(debug_assertions))]
-        type P = Poseidon2GoldilocksHL<POSEIDON2_GOLDILICK_WIDTH>;
+        type P = WP;
+
         type T =
             DuplexChallenger<Self, Self::P, POSEIDON2_GOLDILICK_WIDTH, POSEIDON2_GOLDILICK_RATE>;
         type S = PaddingFreeSponge<Self::P, POSEIDON2_GOLDILICK_WIDTH, POSEIDON2_GOLDILICK_RATE, 4>;
@@ -71,24 +66,34 @@ pub mod impl_goldilocks {
 
         #[cfg(debug_assertions)]
         fn get_default_perm() -> Self::P {
-            Instrumented::new(Poseidon2GoldilocksHL::new(
+            Instrumented::new(Wrapper::new(Poseidon2GoldilocksHL::new(
                 ExternalLayerConstants::<Goldilocks, POSEIDON2_GOLDILICK_WIDTH>::new_from_saved_array(
                     HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS,
                     new_array,
                 ),
                 new_array(HL_GOLDILOCKS_8_INTERNAL_ROUND_CONSTANTS).to_vec(),
-            ))
+            )))
         }
 
         #[cfg(not(debug_assertions))]
         fn get_default_perm() -> Self::P {
-            Poseidon2GoldilocksHL::new(
+            Wrapper::new(Poseidon2GoldilocksHL::new(
                 ExternalLayerConstants::<Goldilocks, POSEIDON2_GOLDILICK_WIDTH>::new_from_saved_array(
                     HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS,
                     new_array,
                 ),
                 new_array(HL_GOLDILOCKS_8_INTERNAL_ROUND_CONSTANTS).to_vec(),
-            )
+            ))
+        }
+
+        fn get_default_perm_rc() -> Vec<Self> {
+            HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS[0]
+                .iter()
+                .flatten()
+                .chain(HL_GOLDILOCKS_8_INTERNAL_ROUND_CONSTANTS.iter())
+                .chain(HL_GOLDILOCKS_8_EXTERNAL_ROUND_CONSTANTS[1].iter().flatten())
+                .map(|v| Self::from_canonical_u64(*v))
+                .collect()
         }
 
         fn get_default_sponge() -> Self::S {

crates/ff_ext/src/lib.rs

@@ -10,6 +10,7 @@ use std::{
     iter::{self, repeat_with},
 };
 mod babybear;
+mod wrapper;
 pub use babybear::impl_babybear::*;
 mod goldilock;
 pub use goldilock::impl_goldilocks::*;

crates/ff_ext/src/poseidon.rs

@@ -2,6 +2,7 @@ use p3::{
     challenger::{FieldChallenger, GrindingChallenger},
     commit::Mmcs,
     field::PrimeField,
+    symmetric::Permutation,
 };
 
 use crate::{ExtensionField, SmallField};
@@ -19,7 +20,7 @@ pub trait FieldChallengerExt<F: PoseidonField>: FieldChallenger<F> {
 
 pub trait PoseidonField: PrimeField + SmallField {
     // permutation
-    type P: Clone;
+    type P: Clone + Permutation<Vec<Self>> + Send + Sync;
     // sponge
     type S: Clone + Sync;
     // compression
@@ -28,6 +29,7 @@ pub trait PoseidonField: PrimeField + SmallField {
     type T: FieldChallenger<Self> + Clone + GrindingChallenger<Witness = Self>;
     fn get_default_challenger() -> Self::T;
     fn get_default_perm() -> Self::P;
+    fn get_default_perm_rc() -> Vec<Self>;
     fn get_default_sponge() -> Self::S;
     fn get_default_compression() -> Self::C;
     fn get_default_mmcs() -> Self::MMCS;
@@ -51,7 +53,7 @@ pub mod impl_instruments {
     };
 
     use once_cell::sync::Lazy;
-    use p3::symmetric::Permutation;
+    use p3::symmetric::{CryptographicPermutation, Permutation};
 
     pub type PermCount = Arc<Mutex<usize>>;
     pub type LabelCounts = Arc<Mutex<HashMap<&'static str, usize>>>;
@@ -136,4 +138,6 @@ pub mod impl_instruments {
             self.inner_perm.permute(input)
         }
     }
+
+    impl<T: Clone, P: Permutation<T>> CryptographicPermutation<T> for Instrumented<P> {}
 }

crates/ff_ext/src/wrapper.rs

@@ -0,0 +1,50 @@
+use p3::symmetric::{CryptographicPermutation, Permutation};
+use std::array::from_fn;
+#[derive(Clone)]
+pub struct Wrapper<P: Clone, const N: usize> {
+    inner: P,
+}
+
+impl<P: Clone, const N: usize> Wrapper<P, N> {
+    pub fn new(inner: P) -> Self {
+        Self { inner }
+    }
+}
+
+impl<const N: usize, T: Clone, P: Permutation<[T; N]> + Clone> Permutation<Vec<T>>
+    for Wrapper<P, N>
+{
+    fn permute(&self, input: Vec<T>) -> Vec<T> {
+        assert_eq!(input.len(), N, "Input vector must be of length {}", N);
+
+        let mut array = from_fn(|i| input[i].clone());
+        self.inner.permute_mut(&mut array);
+
+        array.to_vec()
+    }
+
+    fn permute_mut(&self, input: &mut Vec<T>) {
+        assert_eq!(input.len(), N, "Input vector must be of length {}", N);
+        let mut array = from_fn(|i| input[i].clone());
+        self.inner.permute_mut(&mut array);
+        input.iter_mut().zip(array).for_each(|(i, a)| {
+            *i = a;
+        });
+    }
+}
+impl<const N: usize, T: Clone, P: Permutation<[T; N]> + Clone> Permutation<[T; N]>
+    for Wrapper<P, N>
+{
+    fn permute(&self, input: [T; N]) -> [T; N] {
+        self.inner.permute(input)
+    }
+
+    fn permute_mut(&self, input: &mut [T; N]) {
+        self.inner.permute_mut(input);
+    }
+}
+
+impl<const N: usize, T: Clone, P: Permutation<[T; N]> + Clone> CryptographicPermutation<[T; N]>
+    for Wrapper<P, N>
+{
+}