simplify poseidon permutation calls

Author: TomWambsgans

Cargo.lock

@@ -15,7 +15,6 @@ p3-uni-stark.workspace = true
 p3-matrix.workspace = true
 p3-util.workspace = true
 multilinear-toolkit.workspace = true
-p3-symmetric.workspace = true
 p3-koala-bear.workspace = true
 rand.workspace = true
 whir-p3.workspace = true

crates/air/Cargo.toml

@@ -3,7 +3,6 @@ use multilinear_toolkit::prelude::*;
 use p3_air::BaseAir;
 use p3_field::PrimeField64;
 use p3_koala_bear::{KoalaBear, QuinticExtensionFieldKB};
-use p3_symmetric::Permutation;
 use packed_pcs::{
     ColDims, packed_pcs_commit, packed_pcs_global_statements_for_prover,
     packed_pcs_global_statements_for_verifier, packed_pcs_parse_commitment,
@@ -16,7 +15,7 @@ use utils::{
     MyChallenger, Poseidon16Air, Poseidon24Air, build_poseidon_16_air,
     build_poseidon_16_air_packed, build_poseidon_24_air, build_poseidon_24_air_packed,
     build_prover_state, build_verifier_state, generate_trace_poseidon_16,
-    generate_trace_poseidon_24, get_poseidon16, get_poseidon24, init_tracing,
+    generate_trace_poseidon_24, init_tracing, poseidon16_permute, poseidon24_permute,
 };
 use whir_p3::{
     FoldingFactor, SecurityAssumption, WhirConfig, WhirConfigBuilder, precompute_dft_twiddles,
@@ -111,11 +110,11 @@ fn prepare_poseidon(config: &Poseidon2Config) -> PoseidonSetup {
 
     assert_eq!(
         &witness_matrix_16.values[n_columns_16 - 16..n_columns_16 - 8],
-        &get_poseidon16().permute(witness_matrix_16.values[0..16].try_into().unwrap())[..8]
+        &poseidon16_permute(witness_matrix_16.values[0..16].try_into().unwrap())[..8]
     );
     assert_eq!(
         &witness_matrix_24.values[n_columns_24 - 8..n_columns_24],
-        &get_poseidon24().permute(witness_matrix_24.values[0..24].try_into().unwrap())[16..]
+        &poseidon24_permute(witness_matrix_24.values[0..24].try_into().unwrap())[16..]
     );
 
     let witness_matrix_16_transposed = witness_matrix_16.transpose();

crates/air/src/examples/prove_poseidon2.rs

@@ -1,7 +1,6 @@
 use lean_compiler::*;
 use lean_vm::*;
-use p3_symmetric::Permutation;
-use utils::{get_poseidon16, get_poseidon24};
+use utils::{poseidon16_permute, poseidon24_permute};
 
 const DEFAULT_NO_VEC_RUNTIME_MEMORY: usize = 1 << 15;
 
@@ -327,16 +326,15 @@ fn test_mini_program_3() {
         return;
     }
    "#;
-    let mut public_input: [F; 16] = (0..16).map(F::new).collect::<Vec<F>>().try_into().unwrap();
+    let public_input: [F; 16] = (0..16).map(F::new).collect::<Vec<F>>().try_into().unwrap();
     compile_and_run(
         program.to_string(),
         (&public_input, &[]),
         DEFAULT_NO_VEC_RUNTIME_MEMORY,
         false,
     );
 
-    get_poseidon16().permute_mut(&mut public_input);
-    let _ = public_input;
+    let _ = dbg!(poseidon16_permute(public_input));
 }
 
 #[test]
@@ -355,16 +353,15 @@ fn test_mini_program_4() {
         return;
     }
    "#;
-    let mut public_input: [F; 24] = (0..24).map(F::new).collect::<Vec<F>>().try_into().unwrap();
+    let public_input: [F; 24] = (0..24).map(F::new).collect::<Vec<F>>().try_into().unwrap();
     compile_and_run(
         program.to_string(),
         (&public_input, &[]),
         DEFAULT_NO_VEC_RUNTIME_MEMORY,
         false,
     );
 
-    get_poseidon24().permute_mut(&mut public_input);
-    dbg!(&public_input[16..]);
+    dbg!(&poseidon24_permute(public_input)[16..]);
 }
 
 #[test]

crates/lean_compiler/tests/test_compiler.rs

@@ -13,9 +13,8 @@ use crate::witness::{
 };
 use crate::{CodeAddress, HintExecutionContext, SourceLineNumber};
 use p3_field::PrimeCharacteristicRing;
-use p3_symmetric::Permutation;
 use std::collections::BTreeMap;
-use utils::{get_poseidon16, get_poseidon24, pretty_integer};
+use utils::{poseidon16_permute, poseidon24_permute, pretty_integer};
 use xmss::{Poseidon16History, Poseidon24History};
 
 /// Number of instructions to show in stack trace
@@ -51,10 +50,10 @@ pub fn build_public_memory(public_input: &[F]) -> Vec<F> {
 
     public_memory
         [POSEIDON_16_NULL_HASH_PTR * VECTOR_LEN..(POSEIDON_16_NULL_HASH_PTR + 2) * VECTOR_LEN]
-        .copy_from_slice(&get_poseidon16().permute([F::ZERO; 16]));
+        .copy_from_slice(&poseidon16_permute([F::ZERO; 16]));
     public_memory
         [POSEIDON_24_NULL_HASH_PTR * VECTOR_LEN..(POSEIDON_24_NULL_HASH_PTR + 1) * VECTOR_LEN]
-        .copy_from_slice(&get_poseidon24().permute([F::ZERO; 24])[16..]);
+        .copy_from_slice(&poseidon24_permute([F::ZERO; 24])[16..]);
     public_memory
 }
 

crates/lean_vm/src/execution/runner.rs

@@ -12,10 +12,9 @@ use crate::witness::{
 };
 use multilinear_toolkit::prelude::*;
 use p3_field::{BasedVectorSpace, PrimeCharacteristicRing, dot_product};
-use p3_symmetric::Permutation;
 use p3_util::log2_ceil_usize;
 use std::fmt::{Display, Formatter};
-use utils::{ToUsize, get_poseidon16, get_poseidon24};
+use utils::{ToUsize, poseidon16_permute, poseidon24_permute};
 
 /// Complete set of VM instruction types with comprehensive operation support
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -236,11 +235,7 @@ impl Instruction {
                         *ctx.n_poseidon16_precomputed_used += 1;
                         precomputed.1
                     }
-                    _ => {
-                        let mut output = input;
-                        get_poseidon16().permute_mut(&mut output);
-                        output
-                    }
+                    _ => poseidon16_permute(input),
                 };
 
                 let res0: [F; VECTOR_LEN] = output[..VECTOR_LEN].try_into().unwrap();
@@ -292,8 +287,7 @@ impl Instruction {
                         precomputed.1
                     }
                     _ => {
-                        let mut output = input;
-                        get_poseidon24().permute_mut(&mut output);
+                        let output = poseidon24_permute(input);
                         output[2 * VECTOR_LEN..].try_into().unwrap()
                     }
                 };

crates/lean_vm/src/isa/instruction.rs

@@ -14,6 +14,7 @@ p3-matrix.workspace = true
 p3-challenger.workspace = true
 p3-koala-bear.workspace = true
 tracing-forest.workspace = true
+p3-symmetric.workspace = true
 p3-poseidon2.workspace = true
 p3-poseidon2-air.workspace = true
 tracing-subscriber.workspace = true

crates/utils/Cargo.toml

@@ -19,6 +19,7 @@ use p3_poseidon2_air::p24::Poseidon2Air24;
 use p3_poseidon2_air::p24::RoundConstants24;
 use p3_poseidon2_air::p24::generate_trace_rows_24;
 use p3_poseidon2_air::{p16, p24};
+use p3_symmetric::Permutation;
 
 pub type Poseidon16 = Poseidon2KoalaBear<16>;
 pub type Poseidon24 = Poseidon2KoalaBear<24>;
@@ -81,7 +82,8 @@ pub type MyRoundConstants24<F> = RoundConstants24<F, 24, HALF_FULL_ROUNDS_24, PA
 
 static POSEIDON16_INSTANCE: OnceLock<Poseidon16> = OnceLock::new();
 
-pub fn get_poseidon16() -> &'static Poseidon16 {
+#[inline(always)]
+pub(crate) fn get_poseidon16() -> &'static Poseidon16 {
     POSEIDON16_INSTANCE.get_or_init(|| {
         let round_constants = build_poseidon16_constants();
         let external_constants = ExternalLayerConstants::new(
@@ -95,9 +97,15 @@ pub fn get_poseidon16() -> &'static Poseidon16 {
     })
 }
 
+#[inline(always)]
+pub fn poseidon16_permute(input: [KoalaBear; 16]) -> [KoalaBear; 16] {
+    get_poseidon16().permute(input)
+}
+
 static POSEIDON24_INSTANCE: OnceLock<Poseidon24> = OnceLock::new();
 
-pub fn get_poseidon24() -> &'static Poseidon24 {
+#[inline(always)]
+pub(crate) fn get_poseidon24() -> &'static Poseidon24 {
     POSEIDON24_INSTANCE.get_or_init(|| {
         let round_constants = build_poseidon24_constants();
         let external_constants = ExternalLayerConstants::new(
@@ -111,6 +119,11 @@ pub fn get_poseidon24() -> &'static Poseidon24 {
     })
 }
 
+#[inline(always)]
+pub fn poseidon24_permute(input: [KoalaBear; 24]) -> [KoalaBear; 24] {
+    get_poseidon24().permute(input)
+}
+
 pub fn build_poseidon_16_air() -> Poseidon16Air<KoalaBear> {
     Poseidon16Air::new(build_poseidon16_constants())
 }

crates/utils/src/poseidon2.rs

@@ -11,5 +11,4 @@ p3-koala-bear.workspace = true
 p3-field.workspace = true
 rand.workspace = true
 utils.workspace = true
-p3-symmetric.workspace = true
 p3-util.workspace = true

crates/xmss/Cargo.toml

@@ -1,9 +1,8 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 use p3_koala_bear::KoalaBear;
-use p3_symmetric::Permutation;
-use utils::{get_poseidon16, get_poseidon24};
 
 mod wots;
+use utils::{poseidon16_permute, poseidon24_permute};
 pub use wots::*;
 mod xmss;
 pub use xmss::*;
@@ -24,7 +23,7 @@ pub type Poseidon16History = Vec<([F; 16], [F; 16])>;
 pub type Poseidon24History = Vec<([F; 24], [F; 8])>;
 
 fn poseidon16_compress(a: &Digest, b: &Digest) -> Digest {
-    get_poseidon16().permute([*a, *b].concat().try_into().unwrap())[0..8]
+    poseidon16_permute([*a, *b].concat().try_into().unwrap())[0..8]
         .try_into()
         .unwrap()
 }
@@ -35,13 +34,13 @@ fn poseidon16_compress_with_trace(
     poseidon_16_trace: &mut Vec<([F; 16], [F; 16])>,
 ) -> Digest {
     let input: [F; 16] = [*a, *b].concat().try_into().unwrap();
-    let output = get_poseidon16().permute(input);
+    let output = poseidon16_permute(input);
     poseidon_16_trace.push((input, output));
     output[0..8].try_into().unwrap()
 }
 
 fn poseidon24_compress(a: &Digest, b: &Digest, c: &Digest) -> Digest {
-    get_poseidon24().permute([*a, *b, *c].concat().try_into().unwrap())[16..24]
+    poseidon24_permute([*a, *b, *c].concat().try_into().unwrap())[16..24]
         .try_into()
         .unwrap()
 }
@@ -53,7 +52,7 @@ fn poseidon24_compress_with_trace(
     poseidon_24_trace: &mut Vec<([F; 24], [F; 8])>,
 ) -> Digest {
     let input: [F; 24] = [*a, *b, *c].concat().try_into().unwrap();
-    let output = get_poseidon24().permute(input)[16..24].try_into().unwrap();
+    let output = poseidon24_permute(input)[16..24].try_into().unwrap();
     poseidon_24_trace.push((input, output));
     output
 }