From 2c5847998ee9244d06b97de3c703a573c86c8046 Mon Sep 17 00:00:00 2001 From: wstran Date: Sun, 26 Jul 2026 00:36:35 +0700 Subject: [PATCH 1/3] perf(xmss): stop building the unused WOTS public key when signing xmss_sign spends more time deriving a WOTS public key it never reads than it does producing the signature itself. gen_wots_secret_key goes through WotsSecretKey::new, which walks all V chains to their ends, but the signing path only calls sign_with_encoding, which reads pre_images. The public key is needed by leaf_layer alone, to hash it into a Merkle leaf. Root cause: WotsSecretKey cached `public_key` as a struct field, so every construction paid for the full chain walk whether the caller wanted it or not. Fix: drop the field and compute the key on demand in `public_key(public_param, slot)`. Key generation walks the same chains and produces the same leaves; signing no longer walks them at all. Measured (docker rust:latest, 6 cores, V=42, CHAIN_LENGTH=8, TARGET_SUM=184): chain hashes per signature: 184 useful + 294 dead -> 184 useful gen_wots_secret_key: 370.8 us -> negligible WotsSecretKey is private to the crate (only WotsSignature is re-exported), so this is not a public API change. --- crates/xmss/src/wots.rs | 23 +++++++++++------------ crates/xmss/src/xmss.rs | 10 +++++----- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/crates/xmss/src/wots.rs b/crates/xmss/src/wots.rs index 8fac38378..a811fa8af 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,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( diff --git a/crates/xmss/src/xmss.rs b/crates/xmss/src/xmss.rs index 1fcc3a893..46aae3b7b 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 = 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(); From 05213d3e043f857f722772060f182fb1b0f3b2ea Mon Sep 17 00:00:00 2001 From: wstran Date: Sun, 26 Jul 2026 00:36:51 +0700 Subject: [PATCH 2/3] perf(xmss): make the target-sum encoding allocation-free Signing grinds wots_encode until the codeword sums to TARGET_SUM, which takes on the order of 840 attempts per signature. Every attempt allocated three times over: to_little_endian_bits builds a Vec per encoding field element, flat_map collects a second Vec of 144 bools, and the chunk map collects the codeword Vec. Roughly 839 of 840 attempts discard all of it. Root cause: the codeword was assembled through a bit-vector pipeline instead of being read straight out of the field elements. Fix: read each W-bit chunk with a shift and a mask into a stack array, and accumulate the sum in the same pass. Masking to W bits already bounds every entry below CHAIN_LENGTH, so the target sum is the only remaining condition and is_valid_encoding is no longer needed. Reading chunks per field element rather than from one concatenated bit stream is only equivalent while W divides the bits taken from each element, so that invariant is now named (ENCODING_BITS_PER_FE, CHUNKS_PER_FE) and asserted at compile time, together with CHAIN_LENGTH being a power of two since CHAIN_LENGTH - 1 is used as the chunk mask. Measured (400k attempts per run, 5 interleaved rounds): wots_encode: 2750 ns -> 2534 ns per attempt (-7.9%) The same 502 codewords out of 420000 attempts are accepted before and after, so the encoding is unchanged. --- crates/xmss/src/lib.rs | 11 ++++++++++- crates/xmss/src/wots.rs | 42 +++++++++++++++-------------------------- 2 files changed, 25 insertions(+), 28 deletions(-) diff --git a/crates/xmss/src/lib.rs b/crates/xmss/src/lib.rs index 27de788bd..5189c1ecd 100644 --- a/crates/xmss/src/lib.rs +++ b/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 diff --git a/crates/xmss/src/wots.rs b/crates/xmss/src/wots.rs index a811fa8af..fee003e19 100644 --- a/crates/xmss/src/wots.rs +++ b/crates/xmss/src/wots.rs @@ -133,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::>() - .as_chunks::() - .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::() != 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)] From 390b16b66503b9b723f1c00f62135a1013824072 Mon Sep 17 00:00:00 2001 From: wstran Date: Sun, 26 Jul 2026 00:36:51 +0700 Subject: [PATCH 3/3] test(xmss): add signing benchmarks and signature digests Ignored tests, following the pattern of benchmark_poseidons.rs: signing and verification throughput, per-attempt cost of the grinding loop, and signature digests for fixed (seed, slot, message) triples so that a change meant to leave signatures untouched can be checked against them. --- crates/xmss/src/benchmark_sign.rs | 122 ++++++++++++++++++++++++++++++ 1 file changed, 122 insertions(+) create mode 100644 crates/xmss/src/benchmark_sign.rs diff --git a/crates/xmss/src/benchmark_sign.rs b/crates/xmss/src/benchmark_sign.rs new file mode 100644 index 000000000..7cd9df9ca --- /dev/null +++ b/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}"); + } +}