From 56156862673d6ec12f7489e5194bc840c7291ce7 Mon Sep 17 00:00:00 2001
From: Rick Glenn <therickglenn@gmail.com>
Date: Thu, 30 Jul 2026 17:30:28 -0700
Subject: [PATCH] perf: avoid WOTS public key work during signing

---
 crates/xmss/src/wots.rs | 25 +++++++++++++------------
 crates/xmss/src/xmss.rs | 10 +++++-----
 2 files changed, 18 insertions(+), 17 deletions(-)

diff --git a/crates/xmss/src/wots.rs b/crates/xmss/src/wots.rs
index 8fac3837..f8a41598 100644
--- a/crates/xmss/src/wots.rs
+++ b/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,23 @@ 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
+    /// Derives the public key by walking every chain to its final element.
+    ///
+    /// This is intentionally computed on demand. Signing only needs the secret
+    /// pre-images, so eagerly deriving and storing these chain ends would add
+    /// `V * (CHAIN_LENGTH - 1)` unnecessary hashes to every signature.
+    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(
diff --git a/crates/xmss/src/xmss.rs b/crates/xmss/src/xmss.rs
index 1fcc3a89..46aae3b7 100644
--- a/crates/xmss/src/xmss.rs
+++ b/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();