Speed up signing

benches/benchmark_poseidon.rs

@@ -24,6 +24,8 @@ use leansig::{
     },
 };
 
+type VerificationFixture<T> = (u32, [u8; MESSAGE_LENGTH], T);
+
 /// A template for benchmarking signature schemes (key gen, signing, verification)
 pub fn benchmark_signature_scheme<S: SignatureScheme>(c: &mut Criterion, description: &str) {
     let mut group = c.benchmark_group(format!("Poseidon: {description}"));
@@ -47,7 +49,10 @@ pub fn benchmark_signature_scheme<S: SignatureScheme>(c: &mut Criterion, descrip
 
     group.sample_size(100);
 
-    let (pk, sk) = S::key_gen(&mut rng, 0, S::LIFETIME as usize);
+    // Signing only needs the two prepared bottom trees. Keep the full tree
+    // height, but avoid generating every epoch before a filtered benchmark.
+    let active_epochs = 2 * (S::LIFETIME as usize).isqrt();
+    let (pk, sk) = S::key_gen(&mut rng, 0, active_epochs);
 
     // Get the prepared epoch interval to ensure we only sign at valid epochs
     let prepared_interval = sk.get_prepared_interval();
@@ -62,24 +67,29 @@ pub fn benchmark_signature_scheme<S: SignatureScheme>(c: &mut Criterion, descrip
                 rng.random_range(prepared_interval.start as u32..prepared_interval.end as u32);
 
             // Benchmark signing
-            let _ = S::sign(black_box(&sk), black_box(epoch), black_box(&message));
+            black_box(
+                S::sign(black_box(&sk), black_box(epoch), black_box(&message))
+                    .expect("Signing should succeed"),
+            );
         });
     });
 
-    // Pre-generate messages, epochs, and signatures for verification
-    let precomputed: Vec<(u32, [u8; MESSAGE_LENGTH], S::Signature)> = (0..2000)
-        .map(|_| {
-            let message = rng.random();
-            // Use epochs within the prepared interval
-            let epoch =
-                rng.random_range(prepared_interval.start as u32..prepared_interval.end as u32);
-            let signature = S::sign(&sk, epoch, &message).expect("Signing should succeed");
-            (epoch, message, signature)
-        })
-        .collect();
-
-    // Verification benchmark
+    // Construct verification fixtures only when this benchmark is selected,
+    // outside the timed iterations. A signing-only profile skips this work.
+    let mut verification_fixtures: Option<Vec<VerificationFixture<S::Signature>>> = None;
     group.bench_function("- verify", |b| {
+        let precomputed = verification_fixtures.get_or_insert_with(|| {
+            (0..2000)
+                .map(|_| {
+                    let message = rng.random();
+                    // Use epochs within the prepared interval
+                    let epoch = rng
+                        .random_range(prepared_interval.start as u32..prepared_interval.end as u32);
+                    let signature = S::sign(&sk, epoch, &message).expect("Signing should succeed");
+                    (epoch, message, signature)
+                })
+                .collect()
+        });
         b.iter(|| {
             // Randomly pick a precomputed signature to verify
             let (epoch, message, signature) =

src/inc_encoding.rs

@@ -45,6 +45,23 @@ pub trait IncomparableEncoding {
         randomness: &Self::Randomness,
         epoch: u32,
     ) -> Result<Vec<u8>, Self::Error>;
+
+    /// Find the first valid encoding in counter order. Implementations may
+    /// prepare invariant inputs or batch trials, but must preserve this order.
+    fn find_encoding(
+        parameter: &Self::Parameter,
+        message: &[u8; MESSAGE_LENGTH],
+        epoch: u32,
+        randomness: impl Fn(u64) -> Self::Randomness + Sync,
+    ) -> Option<(Self::Randomness, Vec<u8>)> {
+        for counter in 0..Self::MAX_TRIES {
+            let rho = randomness(counter as u64);
+            if let Ok(chunks) = Self::encode(parameter, message, &rho, epoch) {
+                return Some((rho, chunks));
+            }
+        }
+        None
+    }
 }
 
 pub mod target_sum;

src/inc_encoding/target_sum.rs

@@ -55,6 +55,27 @@ impl<MH: MessageHash, const TARGET_SUM: usize> IncomparableEncoding
         MH::rand(rng)
     }
 
+    fn find_encoding(
+        parameter: &Self::Parameter,
+        message: &[u8; MESSAGE_LENGTH],
+        epoch: u32,
+        randomness: impl Fn(u64) -> Self::Randomness + Sync,
+    ) -> Option<(Self::Randomness, Vec<u8>)> {
+        // Use encode once to retain its compile-time parameter checks.
+        let rho = randomness(0);
+        if let Ok(chunks) = Self::encode(parameter, message, &rho, epoch) {
+            return Some((rho, chunks));
+        }
+        MH::find_target_sum(
+            parameter,
+            message,
+            epoch,
+            TARGET_SUM,
+            Self::MAX_TRIES - 1,
+            |counter| randomness(counter + 1),
+        )
+    }
+
     fn encode(
         parameter: &Self::Parameter,
         message: &[u8; MESSAGE_LENGTH],

src/lib.rs

@@ -23,8 +23,7 @@ pub mod symmetric;
 
 // Cached Poseidon1 permutations.
 //
-// We cache the default Plonky3 Poseidon1 instances once and return a clone.
-// Returning by value preserves existing call sites that take `&perm`.
+// Share immutable round constants instead of cloning their heap-backed vectors.
 
 /// A lazily-initialized, thread-safe cache for the Poseidon1 permutation with a width of 24.
 static POSEIDON1_24: OnceLock<Poseidon1KoalaBear<24>> = OnceLock::new();
@@ -33,15 +32,11 @@ static POSEIDON1_24: OnceLock<Poseidon1KoalaBear<24>> = OnceLock::new();
 static POSEIDON1_16: OnceLock<Poseidon1KoalaBear<16>> = OnceLock::new();
 
 /// Poseidon1 permutation (width 24)
-pub(crate) fn poseidon1_24() -> Poseidon1KoalaBear<24> {
-    POSEIDON1_24
-        .get_or_init(default_koalabear_poseidon1_24)
-        .clone()
+pub(crate) fn poseidon1_24() -> &'static Poseidon1KoalaBear<24> {
+    POSEIDON1_24.get_or_init(default_koalabear_poseidon1_24)
 }
 
 /// Poseidon1 permutation (width 16)
-pub(crate) fn poseidon1_16() -> Poseidon1KoalaBear<16> {
-    POSEIDON1_16
-        .get_or_init(default_koalabear_poseidon1_16)
-        .clone()
+pub(crate) fn poseidon1_16() -> &'static Poseidon1KoalaBear<16> {
+    POSEIDON1_16.get_or_init(default_koalabear_poseidon1_16)
 }

src/signature/generalized_xmss.rs

@@ -876,39 +876,15 @@ where
         let path = combined_path(&sk.top_tree, bottom_tree, epoch);
 
         // now, we need to encode our message using the incomparable encoding.
-        // we retry until we get a valid codeword, or until we give up.
-        let max_tries = IE::MAX_TRIES;
-        let mut attempts = 0;
-        let mut x = None;
-        let mut rho = None;
-        while attempts < max_tries {
-            // get a randomness and try to encode the message. Note: we get the randomness from the PRF
-            // which ensures that signing is deterministic. The PRF is applied to the message and the epoch.
-            // While the intention is that users of the scheme never call sign twice with the same (epoch, sk) pair,
-            // this deterministic approach ensures that calling sign twice is fine, as long as the message stays the same.
-            let curr_rho = PRF::get_randomness(&sk.prf_key, epoch, message, attempts as u64).into();
-            let curr_x = IE::encode(&sk.parameter.into(), message, &curr_rho, epoch);
-
-            // check if we have found a valid codeword, and if so, stop searching
-            if curr_x.is_ok() {
-                rho = Some(curr_rho);
-                x = curr_x.ok();
-                break;
-            }
-
-            attempts += 1;
-        }
-
-        // if we have not found a valid codeword, return an error
-        if x.is_none() {
-            return Err(SigningError::EncodingAttemptsExceeded {
-                attempts: max_tries,
-            });
-        }
-
-        // otherwise, unwrap x and rho
-        let x = x.unwrap();
-        let rho = rho.unwrap();
+        // Search in counter order, allowing the encoding to prepare invariant
+        // inputs and batch retries while retaining deterministic signatures.
+        let randomness = PRF::randomness_generator(&sk.prf_key, epoch, message);
+        let (rho, x) = IE::find_encoding(&sk.parameter.into(), message, epoch, |counter| {
+            randomness(counter).into()
+        })
+        .ok_or(SigningError::EncodingAttemptsExceeded {
+            attempts: IE::MAX_TRIES,
+        })?;
 
         // we will include rho in the signature, and
         // we use x to determine how far the signer walks in the chains
@@ -918,17 +894,7 @@ where
             "Encoding is broken: returned too many or too few chunks."
         );
 
-        // In parallel, compute the hash values for each chain based on the codeword `x`.
-        let hashes = (0..num_chains)
-            .into_par_iter()
-            .map(|chain_index| {
-                // get back to the start of the chain from the PRF
-                let start = PRF::get_domain_element(&sk.prf_key, epoch, chain_index as u64).into();
-                // now walk the chain for a number of steps determined by the current chunk of x
-                let steps = x[chain_index] as usize;
-                chain::<TH>(&sk.parameter, epoch, chain_index as u8, 0, steps, &start)
-            })
-            .collect();
+        let hashes = TH::sign_chains::<PRF>(&sk.prf_key, &sk.parameter, epoch, &x);
 
         // assemble the signature: Merkle path, randomness, chain elements
         Ok(GeneralizedXMSSSignature { path, rho, hashes })
@@ -1094,6 +1060,31 @@ mod tests {
 
     use super::*;
 
+    #[test]
+    fn signing_matches_pre_optimization_fingerprint() {
+        use super::instantiations_poseidon::lifetime_2_to_the_18::target_sum::SIGTargetSumLifetime18W8NoOff as S;
+        use rand::{SeedableRng, rngs::StdRng};
+        use sha3::{Digest, Sha3_256};
+        let mut rng = StdRng::seed_from_u64(0x01ea_0516);
+        let (pk, sk) = S::key_gen(&mut rng, 0, 1024);
+        let inputs: Vec<(u32, [u8; MESSAGE_LENGTH])> = sk
+            .get_prepared_interval()
+            .map(|epoch| (epoch as u32, rng.random()))
+            .collect();
+        let mut digest = Sha3_256::new();
+        for (epoch, message) in inputs.iter().take(16) {
+            let signature = S::sign(&sk, *epoch, message).unwrap();
+            assert!(S::verify(&pk, *epoch, message, &signature));
+            digest.update(signature.to_bytes());
+        }
+        // Captured from the original implementation with this Cargo.lock,
+        // seed, activation interval, and ordered message corpus.
+        assert_eq!(
+            format!("{:x}", digest.finalize()),
+            "ce6c42afc1bb5fa2af80d33e80b583d57d68612c429eda1dff29db545b99242e"
+        );
+    }
+
     use crate::array::FieldArray;
     use p3_field::PrimeField32;
     use proptest::prelude::*;

src/symmetric/message_hash.rs

@@ -37,6 +37,28 @@ pub trait MessageHash {
         randomness: &Self::Randomness,
         message: &[u8; MESSAGE_LENGTH],
     ) -> Result<Vec<u8>, Self::Error>;
+
+    /// Search for the first hash whose chunks have the requested sum.
+    /// The default preserves scalar behavior for hash implementations without
+    /// a prepared or batched search path.
+    fn find_target_sum(
+        parameter: &Self::Parameter,
+        message: &[u8; MESSAGE_LENGTH],
+        epoch: u32,
+        target: usize,
+        max_tries: usize,
+        randomness: impl Fn(u64) -> Self::Randomness + Sync,
+    ) -> Option<(Self::Randomness, Vec<u8>)> {
+        for counter in 0..max_tries {
+            let rho = randomness(counter as u64);
+            if let Ok(chunks) = Self::apply(parameter, epoch, &rho, message) {
+                if chunks.iter().map(|&x| x as usize).sum::<usize>() == target {
+                    return Some((rho, chunks));
+                }
+            }
+        }
+        None
+    }
 }
 
 pub mod aborting;