perf(xmss): remove dead work from the signing path

crates/xmss/src/benchmark_sign.rs

@@ -0,0 +1,122 @@
+use std::hint::black_box;
+use std::time::Instant;
+
+use backend::PrimeCharacteristicRing;
+
+use crate::wots::wots_encode;
+use crate::{
+    CHAIN_LENGTH, F, MESSAGE_LEN_BYTES, Randomness, TARGET_SUM, V, hash_message, xmss_key_gen_from_seed, xmss_sign,
+    xmss_verify,
+};
+
+const SEED: [u8; 32] = [7u8; 32];
+const ACTIVATION_SLOT: u64 = 1 << 20;
+const NUM_ACTIVE_SLOTS: u64 = 1 << 10;
+
+/// Signing and verification throughput.
+///
+/// cargo test --release --package xmss --lib -- benchmark_sign::bench_sign --exact --nocapture --ignored
+#[test]
+#[ignore]
+fn bench_sign() {
+    const WARMUP: u32 = 64;
+    const SAMPLES: u32 = 512;
+
+    let (pk, sk) = xmss_key_gen_from_seed(SEED, ACTIVATION_SLOT, NUM_ACTIVE_SLOTS).unwrap();
+    let message: [u8; MESSAGE_LEN_BYTES] = std::array::from_fn(|i| i as u8);
+
+    for i in 0..WARMUP {
+        let slot = (ACTIVATION_SLOT + u64::from(i)) as u32;
+        let _ = black_box(xmss_sign(&sk, slot, &message).unwrap());
+    }
+
+    let time = Instant::now();
+    for i in 0..SAMPLES {
+        let slot = (ACTIVATION_SLOT + u64::from(WARMUP + i)) as u32;
+        let _ = black_box(xmss_sign(&sk, slot, &message).unwrap());
+    }
+    let sign_elapsed = time.elapsed();
+
+    let slot = (ACTIVATION_SLOT + u64::from(WARMUP)) as u32;
+    let signature = xmss_sign(&sk, slot, &message).unwrap();
+    assert!(xmss_verify(&pk, slot, &message, &signature).is_ok());
+
+    let time = Instant::now();
+    for _ in 0..SAMPLES {
+        let _ = black_box(xmss_verify(&pk, slot, &message, &signature));
+    }
+    let verify_elapsed = time.elapsed();
+
+    println!(
+        "xmss_sign  : {:>9.1} us/op ({:.0} sig/s)",
+        sign_elapsed.as_secs_f64() * 1e6 / f64::from(SAMPLES),
+        f64::from(SAMPLES) / sign_elapsed.as_secs_f64()
+    );
+    println!(
+        "xmss_verify: {:>9.1} us/op",
+        verify_elapsed.as_secs_f64() * 1e6 / f64::from(SAMPLES)
+    );
+    println!(
+        "chain hashes: {TARGET_SUM} walked per signature, {} in a full WOTS public key",
+        V * (CHAIN_LENGTH - 1)
+    );
+}
+
+/// Per-attempt cost of the loop that grinds randomness until the codeword hits `TARGET_SUM`.
+/// The accepted-codeword count doubles as a check that the encoding itself did not change.
+///
+/// cargo test --release --package xmss --lib -- benchmark_sign::bench_wots_encode --exact --nocapture --ignored
+#[test]
+#[ignore]
+fn bench_wots_encode() {
+    const WARMUP: u64 = 20_000;
+    const ATTEMPTS: u64 = 400_000;
+
+    let (pk, _) = xmss_key_gen_from_seed(SEED, ACTIVATION_SLOT, NUM_ACTIVE_SLOTS).unwrap();
+    let message: [u8; MESSAGE_LEN_BYTES] = std::array::from_fn(|i| i as u8);
+    let message_fe = hash_message(&message);
+    let slot = ACTIVATION_SLOT as u32;
+    // Stand-in for the per-attempt PRF output: only its spread matters here.
+    let randomness = |n: u64| -> Randomness {
+        std::array::from_fn(|i| F::from_usize((n.wrapping_mul(0x9e37_79b9) as usize).wrapping_add(i)))
+    };
+
+    let mut accepted = 0u64;
+    for n in 0..WARMUP {
+        accepted += u64::from(black_box(wots_encode(&message_fe, slot, &pk, &randomness(n))).is_some());
+    }
+
+    let time = Instant::now();
+    for n in WARMUP..WARMUP + ATTEMPTS {
+        accepted += u64::from(black_box(wots_encode(&message_fe, slot, &pk, &randomness(n))).is_some());
+    }
+    let elapsed = time.elapsed();
+
+    println!(
+        "wots_encode: {:>8.1} ns/attempt ({accepted} accepted of {})",
+        elapsed.as_secs_f64() * 1e9 / ATTEMPTS as f64,
+        WARMUP + ATTEMPTS
+    );
+}
+
+/// Signature digests for fixed (seed, slot, message) triples. An optimization that is meant to
+/// leave the signature untouched must print the same digests before and after.
+///
+/// cargo test --release --package xmss --lib -- benchmark_sign::signature_digests --exact --nocapture --ignored
+#[test]
+#[ignore]
+fn signature_digests() {
+    let (pk, sk) = xmss_key_gen_from_seed(SEED, ACTIVATION_SLOT, NUM_ACTIVE_SLOTS).unwrap();
+    for slot_offset in [0u64, 1, 7, 511, 1023] {
+        let slot = (ACTIVATION_SLOT + slot_offset) as u32;
+        let message: [u8; MESSAGE_LEN_BYTES] = std::array::from_fn(|i| (i as u8).wrapping_mul(slot as u8));
+        let signature = xmss_sign(&sk, slot, &message).unwrap();
+        assert!(xmss_verify(&pk, slot, &message, &signature).is_ok());
+
+        let mut digest = 0xcbf2_9ce4_8422_2325u64;
+        for byte in ssz::Encode::as_ssz_bytes(&signature) {
+            digest = (digest ^ u64::from(byte)).wrapping_mul(0x0100_0000_01b3);
+        }
+        println!("slot {slot}: signature fnv1a = {digest:016x}");
+    }
+}

crates/xmss/src/lib.rs

@@ -1,6 +1,8 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 use backend::{DIGEST_LEN_FE, KoalaBear, POSEIDON1_WIDTH, PrimeCharacteristicRing, PrimeField32, poseidon16_compress};
 
+#[cfg(test)]
+mod benchmark_sign;
 #[cfg(any(test, feature = "test-utils"))]
 pub mod signers_cache;
 mod ssz_serialization;
@@ -30,7 +32,12 @@ pub const W: usize = 3;
 pub const CHAIN_LENGTH: usize = 1 << W;
 pub const NUM_CHAIN_HASHES: usize = 110;
 pub const TARGET_SUM: usize = V * (CHAIN_LENGTH - 1) - NUM_CHAIN_HASHES;
-pub const NUM_ENCODING_FE: usize = V.div_ceil(24 / W);
+/// Bits taken from each encoding field element. Below the 31-bit modulus, and a multiple of `W`
+/// so that a chunk never straddles two field elements.
+pub const ENCODING_BITS_PER_FE: usize = 24;
+/// `W`-bit chunks carried by one encoding field element.
+pub const CHUNKS_PER_FE: usize = ENCODING_BITS_PER_FE / W;
+pub const NUM_ENCODING_FE: usize = V.div_ceil(CHUNKS_PER_FE);
 pub const RANDOMNESS_LEN_FE: usize = 6;
 /// Byte length of the messages being signed.
 pub const MESSAGE_LEN_BYTES: usize = 32;
@@ -52,6 +59,8 @@ pub const TWEAK_TYPE_MERKLE: usize = 2;
 pub const TWEAK_TYPE_ENCODING: usize = 3;
 
 const _: () = assert!(V.is_multiple_of(2)); // For efficiency of the snark (we can batch chains in pairs)
+const _: () = assert!(ENCODING_BITS_PER_FE.is_multiple_of(W)); // A codeword chunk never straddles two field elements
+const _: () = assert!(CHAIN_LENGTH.is_power_of_two()); // CHAIN_LENGTH - 1 is the chunk bit mask
 const _: () = assert!(MESSAGE_EMBEDDING_LEN_FE * 30 >= MESSAGE_LEN_BYTES * 8); // Injective embedding
 const _: () = assert!(MESSAGE_LEN_FE == DIGEST_LEN_FE); // hash_message output is one Poseidon digest
 const _: () = assert!(MESSAGE_EMBEDDING_LEN_FE < POSEIDON1_WIDTH); // Domain sep + embedding fit

crates/xmss/src/wots.rs

@@ -7,7 +7,6 @@ use crate::*;
 /// No Debug: `pre_images` are the one-time secret keys.
 pub struct WotsSecretKey {
     pre_images: [Digest; V],
-    public_key: WotsPublicKey,
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -24,21 +23,21 @@ pub struct WotsSignature {
 }
 
 impl WotsSecretKey {
-    pub fn random(rng: &mut impl CryptoRng, public_param: PublicParam, slot: u32) -> Self {
-        Self::new(rng.random(), public_param, slot)
+    pub fn random(rng: &mut impl CryptoRng) -> Self {
+        Self::new(rng.random())
     }
 
-    pub fn new(pre_images: [Digest; V], public_param: PublicParam, slot: u32) -> Self {
-        Self {
-            pre_images,
-            public_key: WotsPublicKey(std::array::from_fn(|i| {
-                iterate_hash(&pre_images[i], CHAIN_LENGTH - 1, public_param, slot, i, 0)
-            })),
-        }
+    pub const fn new(pre_images: [Digest; V]) -> Self {
+        Self { pre_images }
     }
 
-    pub const fn public_key(&self) -> &WotsPublicKey {
-        &self.public_key
+    /// Walks all `V` chains to their ends. Only key generation needs this (it hashes the WOTS
+    /// public key into a Merkle leaf); signing walks each chain part-way instead, so it must not
+    /// pay for the full walk.
+    pub fn public_key(&self, public_param: PublicParam, slot: u32) -> WotsPublicKey {
+        WotsPublicKey(std::array::from_fn(|i| {
+            iterate_hash(&self.pre_images[i], CHAIN_LENGTH - 1, public_param, slot, i, 0)
+        }))
     }
 
     pub(crate) fn sign_with_encoding(
@@ -134,35 +133,23 @@ pub fn wots_encode(
         // ensures uniformity of encoding
         return None;
     }
-    let all_indices: Vec<_> = compressed[..NUM_ENCODING_FE]
-        .iter()
-        .flat_map(|kb| to_little_endian_bits(kb.to_usize(), 24))
-        .collect::<Vec<_>>()
-        .as_chunks::<W>()
-        .0
-        .iter()
-        .take(V)
-        .map(|chunk| {
-            chunk
-                .iter()
-                .enumerate()
-                .fold(0u8, |acc, (i, &bit)| acc | (u8::from(bit) << i))
-        })
-        .collect();
-    is_valid_encoding(&all_indices).then(|| all_indices[..V].try_into().unwrap())
-}
-
-fn is_valid_encoding(encoding: &[u8]) -> bool {
-    if encoding.len() != V {
-        return false;
-    }
-    if !encoding.iter().all(|&x| (x as usize) < CHAIN_LENGTH) {
-        return false;
+    // Signing grinds this function until the encoding hits `TARGET_SUM`, so it runs on the
+    // order of a thousand times per signature: keep it allocation-free.
+    let mut words = [0usize; NUM_ENCODING_FE];
+    for (word, &kb) in words.iter_mut().zip(&compressed[..NUM_ENCODING_FE]) {
+        *word = kb.to_usize();
     }
-    if encoding.iter().map(|&x| x as usize).sum::<usize>() != TARGET_SUM {
-        return false;
+    let mut encoding = [0u8; V];
+    let mut sum = 0usize;
+    for (i, out) in encoding.iter_mut().enumerate() {
+        // Chunk i lives in word i / CHUNKS_PER_FE, at bit offset W * (i % CHUNKS_PER_FE).
+        let chunk = ((words[i / CHUNKS_PER_FE] >> (W * (i % CHUNKS_PER_FE))) & (CHAIN_LENGTH - 1)) as u8;
+        *out = chunk;
+        sum += usize::from(chunk);
     }
-    true
+    // Masking to W bits already guarantees every entry is below CHAIN_LENGTH, so the target sum
+    // is the only remaining validity condition.
+    (sum == TARGET_SUM).then_some(encoding)
 }
 
 #[cfg(test)]

crates/xmss/src/xmss.rs

@@ -96,14 +96,14 @@ impl XmssPublicKey {
     }
 }
 
-fn gen_wots_secret_key(seed: &[u8; 32], slot: u32, public_param: PublicParam) -> WotsSecretKey {
+fn gen_wots_secret_key(seed: &[u8; 32], slot: u32) -> WotsSecretKey {
     let rng_seed_fe = poseidon_prf(PRF_DOMAINSEP_WOTS_SECRET_KEY, seed, [slot as usize, 0]);
     let mut rng_seed = [0u8; 32];
     for (chunk, f) in rng_seed.as_chunks_mut::<4>().0.iter_mut().zip(rng_seed_fe) {
         *chunk = f.as_canonical_u32().to_le_bytes();
     }
     let mut rng = StdRng::from_seed(rng_seed);
-    WotsSecretKey::random(&mut rng, public_param, slot)
+    WotsSecretKey::random(&mut rng)
 }
 
 fn gen_public_param(seed: &[u8; 32]) -> PublicParam {
@@ -150,8 +150,8 @@ fn leaf_layer(seed: &[u8; 32], public_param: &PublicParam, lo: u64, hi: u64, seq
     let mut leaves: Vec<Digest> = unsafe { uninitialized_vec((hi - lo + 1) as usize) };
     fill(sequential, &mut leaves, |k, out| {
         let slot = (lo + k as u64) as u32;
-        let wots = gen_wots_secret_key(seed, slot, *public_param);
-        *out = wots.public_key().hash(*public_param, slot);
+        let wots = gen_wots_secret_key(seed, slot);
+        *out = wots.public_key(*public_param, slot).hash(*public_param, slot);
     });
     leaves
 }
@@ -353,7 +353,7 @@ pub fn xmss_sign(
             wots_encode(&message_fe, slot, &pub_key, &randomness).map(|encoding| (randomness, encoding))
         })
         .ok_or(XmssSignatureError::EncodingAttemptsExceeded)?;
-    let wots_secret_key = gen_wots_secret_key(&secret_key.seed, slot, secret_key.public_param);
+    let wots_secret_key = gen_wots_secret_key(&secret_key.seed, slot);
     let wots_signature = wots_secret_key.sign_with_encoding(randomness, &encoding, secret_key.public_param, slot);
     let cache = secret_key.cached_bottom_subtree(slot);
     let sub = cache.as_ref().unwrap();