From 6a50096245c729528ea227857869974a3166dd79 Mon Sep 17 00:00:00 2001 From: Anudit Nagar Date: Thu, 1 Oct 2026 10:08:54 +0530 Subject: [PATCH] Speed up signing --- benches/benchmark_poseidon.rs | 40 ++- src/inc_encoding.rs | 17 ++ src/inc_encoding/target_sum.rs | 21 ++ src/lib.rs | 15 +- src/signature/generalized_xmss.rs | 79 +++-- src/symmetric/message_hash.rs | 22 ++ src/symmetric/message_hash/poseidon.rs | 384 ++++++++++++++++++------- src/symmetric/prf.rs | 10 + src/symmetric/prf/shake_to_field.rs | 49 +++- src/symmetric/tweak_hash.rs | 22 ++ src/symmetric/tweak_hash/poseidon.rs | 263 ++++++++++++----- 11 files changed, 675 insertions(+), 247 deletions(-) diff --git a/benches/benchmark_poseidon.rs b/benches/benchmark_poseidon.rs index da21476..309357b 100644 --- a/benches/benchmark_poseidon.rs +++ b/benches/benchmark_poseidon.rs @@ -24,6 +24,8 @@ use leansig::{ }, }; +type VerificationFixture = (u32, [u8; MESSAGE_LENGTH], T); + /// A template for benchmarking signature schemes (key gen, signing, verification) pub fn benchmark_signature_scheme(c: &mut Criterion, description: &str) { let mut group = c.benchmark_group(format!("Poseidon: {description}")); @@ -47,7 +49,10 @@ pub fn benchmark_signature_scheme(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(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>> = 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) = diff --git a/src/inc_encoding.rs b/src/inc_encoding.rs index f5803e2..9f5fa3d 100644 --- a/src/inc_encoding.rs +++ b/src/inc_encoding.rs @@ -45,6 +45,23 @@ pub trait IncomparableEncoding { randomness: &Self::Randomness, epoch: u32, ) -> Result, 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)> { + 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; diff --git a/src/inc_encoding/target_sum.rs b/src/inc_encoding/target_sum.rs index 50b1930..4ad3f67 100644 --- a/src/inc_encoding/target_sum.rs +++ b/src/inc_encoding/target_sum.rs @@ -55,6 +55,27 @@ impl 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)> { + // 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], diff --git a/src/lib.rs b/src/lib.rs index 121c46b..008a144 100644 --- a/src/lib.rs +++ b/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> = OnceLock::new(); @@ -33,15 +32,11 @@ static POSEIDON1_24: OnceLock> = OnceLock::new(); static POSEIDON1_16: OnceLock> = 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) } diff --git a/src/signature/generalized_xmss.rs b/src/signature/generalized_xmss.rs index acb9e2a..5fdbaa9 100644 --- a/src/signature/generalized_xmss.rs +++ b/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::(&sk.parameter, epoch, chain_index as u8, 0, steps, &start) - }) - .collect(); + let hashes = TH::sign_chains::(&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::*; diff --git a/src/symmetric/message_hash.rs b/src/symmetric/message_hash.rs index c50bc8b..844b834 100644 --- a/src/symmetric/message_hash.rs +++ b/src/symmetric/message_hash.rs @@ -37,6 +37,28 @@ pub trait MessageHash { randomness: &Self::Randomness, message: &[u8; MESSAGE_LENGTH], ) -> Result, 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)> { + 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::() == target { + return Some((rho, chunks)); + } + } + } + None + } } pub mod aborting; diff --git a/src/symmetric/message_hash/poseidon.rs b/src/symmetric/message_hash/poseidon.rs index 32c1ae8..4990239 100644 --- a/src/symmetric/message_hash/poseidon.rs +++ b/src/symmetric/message_hash/poseidon.rs @@ -1,14 +1,17 @@ use std::convert::Infallible; use num_bigint::BigUint; +use p3_field::PackedValue; use p3_field::PrimeCharacteristicRing; use p3_field::PrimeField; use p3_field::PrimeField64; +use rayon::prelude::*; use serde::{Serialize, de::DeserializeOwned}; use super::MessageHash; use crate::F; use crate::MESSAGE_LENGTH; +use crate::PackedF; use crate::TWEAK_SEPARATOR_FOR_MESSAGE_HASH; use crate::array::FieldArray; use crate::poseidon1_24; @@ -78,6 +81,31 @@ pub fn encode_epoch(epoch: u32) -> [F; TWEAK_LEN_FE] fn decode_to_chunks( field_elements: &[F; HASH_LEN_FE], ) -> [u8; DIMENSION] { + if BASE.is_power_of_two() { + // HASH_LEN_FE base-p digits fit in HASH_LEN_FE 32-bit limbs since + // KoalaBear is a 31-bit field. Keep the original big-endian digit order. + let mut limbs = [0u32; HASH_LEN_FE]; + for fe in field_elements { + let mut carry = fe.as_canonical_u64(); + for limb in &mut limbs { + let value = (*limb as u64) * F::ORDER_U64 + carry; + *limb = value as u32; + carry = value >> 32; + } + debug_assert_eq!(carry, 0); + } + let bits = BASE.trailing_zeros() as usize; + return std::array::from_fn(|i| { + let bit = i * bits; + let index = bit / 32; + let shift = bit % 32; + let mut value = limbs.get(index).copied().unwrap_or(0) as u64 >> shift; + if shift + bits > 32 { + value |= (limbs.get(index + 1).copied().unwrap_or(0) as u64) << (32 - shift); + } + (value & (BASE as u64 - 1)) as u8 + }); + } // Combine field elements into one big integer let mut acc = BigUint::ZERO; for fe in field_elements { @@ -122,7 +150,7 @@ pub(crate) fn poseidon_message_hash_fe< .copied() .collect(); - poseidon_compress::(&perm, &combined_input_vec) + poseidon_compress::(perm, &combined_input_vec) } /// A message hash implemented using Poseidon1 @@ -144,6 +172,126 @@ pub struct PoseidonMessageHash< const MSG_LEN_FE: usize, >; +impl< + const PARAMETER_LEN: usize, + const RAND_LEN_FE: usize, + const HASH_LEN_FE: usize, + const DIMENSION: usize, + const BASE: usize, + const TWEAK_LEN_FE: usize, + const MSG_LEN_FE: usize, +> + PoseidonMessageHash< + PARAMETER_LEN, + RAND_LEN_FE, + HASH_LEN_FE, + DIMENSION, + BASE, + TWEAK_LEN_FE, + MSG_LEN_FE, + > +{ + // Compile-time parameter validation for Poseidon Message Hash + // + // This hash implements Th_msg from DKKW25. It takes: + // + // Th_msg(P, T, M, R) = Decode_{p,η',w}( + // PoseidonCompress_{p, t_msg, η'}(R || P || EncT(T) || EncM(M)) + // ) + // + // The output is a vector of DIMENSION chunks, each in {0, ..., BASE-1}. + // This is used inside the incomparable encoding to map (message, seed) + // to a hypercube vertex x ∈ Z_w^v before the target sum filter. + // + // DKKW25: https://eprint.iacr.org/2025/055 + const PARAMETER_CHECK: () = { + // Poseidon capacity constraints + // + // We use Poseidon in compression mode with a width-24 permutation. + // The entire input (R || P || T || M) must fit in one call, and + // the output η' field elements are extracted from the same state. + assert!( + PARAMETER_LEN + TWEAK_LEN_FE + RAND_LEN_FE + MSG_LEN_FE <= 24, + "Poseidon of width 24 is not enough for the input" + ); + assert!( + HASH_LEN_FE <= 24, + "Poseidon of width 24 is not enough for the output" + ); + + // Non-trivial output + // + // The hash must produce at least one field element to be useful. + assert!( + HASH_LEN_FE >= 1, + "Poseidon Message Hash: HASH_LEN_FE must be at least 1" + ); + + // Representation constraints + // + // The Generalized XMSS construction (DKKW25) stores chunk values + // and chain indices as u8 in the signature and tweak encodings. + // BASE (= w) and DIMENSION (= v) must each fit in one byte. + assert!( + BASE <= 1 << 8, + "Poseidon Message Hash: Base must be at most 2^8" + ); + assert!( + DIMENSION <= 1 << 8, + "Poseidon Message Hash: Dimension must be at most 2^8" + ); + + // Encoding well-formedness + // + // Definition 13 (DKKW25): the incomparable encoding maps into + // {0, ..., w-1}^v with w ≥ 2. A single-element alphabet makes + // the code trivial and incomparability vacuous. + assert!(BASE >= 2, "Poseidon Message Hash: BASE must be at least 2"); + + // Injective encoding of inputs + // + // The message (32 bytes = 256 bits) is encoded into MSG_LEN_FE + // field elements via base-p decomposition. For this to be + // injective, the output space must be at least as large: + // + // p^MSG_LEN_FE ≥ 2^(8 * MESSAGE_LENGTH) + // ⟹ ⌊log2(p)⌋ * MSG_LEN_FE ≥ 8 * MESSAGE_LENGTH + let bits_per_fe = F::ORDER_U64.ilog2() as usize; + assert!( + bits_per_fe * MSG_LEN_FE >= 8 * MESSAGE_LENGTH, + "Poseidon Message Hash: not enough field elements to encode the message" + ); + + // The epoch tweak (DKKW25) packs a u32 epoch and an + // 8-bit domain separator into one value, requiring 40 bits total: + // + // tweakm(ep) = (0x02 || ep) ← 8 + 32 = 40 bits + // sep epoch + // + // The field element representation needs enough capacity: + // ⌊log2(p)⌋ * TWEAK_LEN_FE ≥ 40 + assert!( + bits_per_fe * TWEAK_LEN_FE >= 40, + "Poseidon Message Hash: not enough field elements to encode the epoch tweak" + ); + + // Injective decoding to chunks + // + // The hash output (η' field elements) is decoded via Decode_{p,η',w} + // (DKKW25): interpret the η' elements as one big integer + // and express it in base w to get DIMENSION chunks. + // + // For injectivity, the big integer must have enough room: + // p^η' ≤ w^DIMENSION + // ⟹ ⌊log2(p)⌋ * HASH_LEN_FE ≤ DIMENSION * ⌈log2(w)⌉ + let chunk_size = (usize::BITS - (BASE - 1).leading_zeros()) as usize; + assert!( + bits_per_fe * HASH_LEN_FE <= DIMENSION * chunk_size, + "Poseidon Message Hash: not enough chunks to decode the hash output" + ); + }; +} + impl< const PARAMETER_LEN: usize, const RAND_LEN_FE: usize, @@ -180,111 +328,84 @@ where FieldArray(rng.random()) } + 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)> { + let () = Self::PARAMETER_CHECK; + let message_fe = encode_message::(message); + let epoch_fe = encode_epoch::(epoch); + let mut input = [PackedF::ZERO; 24]; + for (dest, value) in input.iter_mut().zip( + message_fe + .iter() + .chain(parameter.iter()) + .chain(epoch_fe.iter()), + ) { + *dest = PackedF::from(*value); + } + let rand_offset = MSG_LEN_FE + PARAMETER_LEN + TWEAK_LEN_FE; + let search = |start: usize, end: usize| { + let mut input = input; + for first in (start..end).step_by(PackedF::WIDTH) { + let count = (end - first).min(PackedF::WIDTH); + let mut rhos = [FieldArray([F::ZERO; RAND_LEN_FE]); PackedF::WIDTH]; + for (lane, rho) in rhos.iter_mut().take(count).enumerate() { + *rho = randomness((first + lane) as u64); + } + for i in 0..RAND_LEN_FE { + input[rand_offset + i] = PackedF::from_fn(|lane| rhos[lane % count][i]); + } + let output = + poseidon_compress::(poseidon1_24(), &input); + for (lane, rho) in rhos.into_iter().take(count).enumerate() { + let scalar = std::array::from_fn(|i| output[i].as_slice()[lane]); + let chunks = decode_to_chunks::(&scalar); + if chunks.iter().map(|&x| x as usize).sum::() == target { + return Some((rho, chunks.to_vec())); + } + } + } + None + }; + if BASE == 256 { + // Bound speculative work to one window. find_first waits for every + // earlier block, preserving the scalar search's first valid counter. + const WINDOW: usize = 1024; + const BLOCK: usize = 64; + for start in (0..max_tries).step_by(WINDOW) { + let end = (start + WINDOW).min(max_tries); + if let Some(result) = (0..(end - start).div_ceil(BLOCK)) + .into_par_iter() + .map(|block| { + search( + start + block * BLOCK, + (start + (block + 1) * BLOCK).min(end), + ) + }) + .find_first(Option::is_some) + .flatten() + { + return Some(result); + } + } + None + } else { + search(0, max_tries) + } + } + fn apply( parameter: &Self::Parameter, epoch: u32, randomness: &Self::Randomness, message: &[u8; MESSAGE_LENGTH], ) -> Result, Infallible> { - // Compile-time parameter validation for Poseidon Message Hash - // - // This hash implements Th_msg from DKKW25. It takes: - // - // Th_msg(P, T, M, R) = Decode_{p,η',w}( - // PoseidonCompress_{p, t_msg, η'}(R || P || EncT(T) || EncM(M)) - // ) - // - // The output is a vector of DIMENSION chunks, each in {0, ..., BASE-1}. - // This is used inside the incomparable encoding to map (message, seed) - // to a hypercube vertex x ∈ Z_w^v before the target sum filter. - // - // DKKW25: https://eprint.iacr.org/2025/055 - const { - // Poseidon capacity constraints - // - // We use Poseidon in compression mode with a width-24 permutation. - // The entire input (R || P || T || M) must fit in one call, and - // the output η' field elements are extracted from the same state. - assert!( - PARAMETER_LEN + TWEAK_LEN_FE + RAND_LEN_FE + MSG_LEN_FE <= 24, - "Poseidon of width 24 is not enough for the input" - ); - assert!( - HASH_LEN_FE <= 24, - "Poseidon of width 24 is not enough for the output" - ); - - // Non-trivial output - // - // The hash must produce at least one field element to be useful. - assert!( - HASH_LEN_FE >= 1, - "Poseidon Message Hash: HASH_LEN_FE must be at least 1" - ); - - // Representation constraints - // - // The Generalized XMSS construction (DKKW25) stores chunk values - // and chain indices as u8 in the signature and tweak encodings. - // BASE (= w) and DIMENSION (= v) must each fit in one byte. - assert!( - BASE <= 1 << 8, - "Poseidon Message Hash: Base must be at most 2^8" - ); - assert!( - DIMENSION <= 1 << 8, - "Poseidon Message Hash: Dimension must be at most 2^8" - ); - - // Encoding well-formedness - // - // Definition 13 (DKKW25): the incomparable encoding maps into - // {0, ..., w-1}^v with w ≥ 2. A single-element alphabet makes - // the code trivial and incomparability vacuous. - assert!(BASE >= 2, "Poseidon Message Hash: BASE must be at least 2"); - - // Injective encoding of inputs - // - // The message (32 bytes = 256 bits) is encoded into MSG_LEN_FE - // field elements via base-p decomposition. For this to be - // injective, the output space must be at least as large: - // - // p^MSG_LEN_FE ≥ 2^(8 * MESSAGE_LENGTH) - // ⟹ ⌊log2(p)⌋ * MSG_LEN_FE ≥ 8 * MESSAGE_LENGTH - let bits_per_fe = F::ORDER_U64.ilog2() as usize; - assert!( - bits_per_fe * MSG_LEN_FE >= 8 * MESSAGE_LENGTH, - "Poseidon Message Hash: not enough field elements to encode the message" - ); - - // The epoch tweak (DKKW25) packs a u32 epoch and an - // 8-bit domain separator into one value, requiring 40 bits total: - // - // tweakm(ep) = (0x02 || ep) ← 8 + 32 = 40 bits - // sep epoch - // - // The field element representation needs enough capacity: - // ⌊log2(p)⌋ * TWEAK_LEN_FE ≥ 40 - assert!( - bits_per_fe * TWEAK_LEN_FE >= 40, - "Poseidon Message Hash: not enough field elements to encode the epoch tweak" - ); - - // Injective decoding to chunks - // - // The hash output (η' field elements) is decoded via Decode_{p,η',w} - // (DKKW25): interpret the η' elements as one big integer - // and express it in base w to get DIMENSION chunks. - // - // For injectivity, the big integer must have enough room: - // p^η' ≤ w^DIMENSION - // ⟹ ⌊log2(p)⌋ * HASH_LEN_FE ≤ DIMENSION * ⌈log2(w)⌉ - let chunk_size = (usize::BITS - (BASE - 1).leading_zeros()) as usize; - assert!( - bits_per_fe * HASH_LEN_FE <= DIMENSION * chunk_size, - "Poseidon Message Hash: not enough chunks to decode the hash output" - ); - } + let () = Self::PARAMETER_CHECK; let hash_fe = poseidon_message_hash_fe::< PARAMETER_LEN, @@ -313,6 +434,69 @@ mod tests { use rand::RngExt; use std::collections::HashMap; + fn reference_chunks(input: &[F; 5]) -> [u8; D] { + let mut acc = BigUint::ZERO; + for fe in input { + acc = acc * F::ORDER_U64 + fe.as_canonical_biguint(); + } + std::array::from_fn(|_| { + let digit = (&acc % B).try_into().unwrap(); + acc /= B; + digit + }) + } + + proptest! { + #[test] + fn fixed_limbs_match_biguint(values in prop::array::uniform5(0u32..F::ORDER_U32)) { + let input = values.map(F::new); + // Include chunks beyond the integer's bit length and widths which + // cross limb boundaries, plus the non-power-of-two fallback. + prop_assert_eq!(decode_to_chunks::<163, 2, 5>(&input), reference_chunks::<163, 2>(&input)); + prop_assert_eq!(decode_to_chunks::<78, 4, 5>(&input), reference_chunks::<78, 4>(&input)); + prop_assert_eq!(decode_to_chunks::<39, 16, 5>(&input), reference_chunks::<39, 16>(&input)); + prop_assert_eq!(decode_to_chunks::<20, 256, 5>(&input), reference_chunks::<20, 256>(&input)); + prop_assert_eq!(decode_to_chunks::<60, 8, 5>(&input), reference_chunks::<60, 8>(&input)); + prop_assert_eq!(decode_to_chunks::<40, 32, 5>(&input), reference_chunks::<40, 32>(&input)); + prop_assert_eq!(decode_to_chunks::<80, 3, 5>(&input), reference_chunks::<80, 3>(&input)); + } + } + + #[test] + fn batched_search_matches_first_scalar_counter_and_limit() { + use crate::symmetric::prf::{Pseudorandom, shake_to_field::ShakePRFtoF}; + use rand::{SeedableRng, rngs::StdRng}; + type MH = PoseidonMessageHash<5, 6, 5, 20, 256, 2, 9>; + type PRF = ShakePRFtoF<7, 6>; + let mut rng = StdRng::seed_from_u64(42); + let parameter = FieldArray(rng.random::<[F; 5]>()); + let key = PRF::key_gen(&mut rng); + for epoch in 0..4 { + let message = rng.random(); + let randomness = + |counter| FieldArray(PRF::get_randomness(&key, epoch, &message, counter)); + for limit in [0, 1, 3, 4, 5, 127, 128, 129, 1023, 1024, 1025, 4096] { + let expected = (0..limit).find_map(|counter| { + let rho = randomness(counter as u64); + let chunks = MH::apply(¶meter, epoch, &rho, &message).unwrap(); + (chunks.iter().map(|&x| x as usize).sum::() == 2550) + .then_some((rho, chunks)) + }); + assert_eq!( + MH::find_target_sum(¶meter, &message, epoch, 2550, limit, &randomness), + expected + ); + } + // Counter zero must also be eligible, including a one-trial limit. + let chunks = MH::apply(¶meter, epoch, &randomness(0), &message).unwrap(); + let target = chunks.iter().map(|&x| x as usize).sum(); + assert_eq!( + MH::find_target_sum(¶meter, &message, epoch, target, 1, &randomness), + Some((randomness(0), chunks)) + ); + } + } + #[test] fn test_apply() { let mut rng = rand::rng(); diff --git a/src/symmetric/prf.rs b/src/symmetric/prf.rs index edc7c9f..9072a03 100644 --- a/src/symmetric/prf.rs +++ b/src/symmetric/prf.rs @@ -25,6 +25,16 @@ pub trait Pseudorandom { message: &[u8; MESSAGE_LENGTH], counter: u64, ) -> Self::Randomness; + + /// Prepare the invariant prefix of a randomness stream. The returned + /// function must match get_randomness for every counter. + fn randomness_generator<'a>( + key: &'a Self::Key, + epoch: u32, + message: &'a [u8; MESSAGE_LENGTH], + ) -> impl Fn(u64) -> Self::Randomness + Sync + 'a { + move |counter| Self::get_randomness(key, epoch, message, counter) + } } pub mod shake_to_field; diff --git a/src/symmetric/prf/shake_to_field.rs b/src/symmetric/prf/shake_to_field.rs index 28c794f..6fe373c 100644 --- a/src/symmetric/prf/shake_to_field.rs +++ b/src/symmetric/prf/shake_to_field.rs @@ -1,7 +1,7 @@ use crate::F; use super::Pseudorandom; -use p3_field::PrimeCharacteristicRing; +use p3_field::{PrimeCharacteristicRing, PrimeField64}; use serde::{Serialize, de::DeserializeOwned}; use sha3::{ Shake128, @@ -18,6 +18,12 @@ const PRF_DOMAIN_SEP: [u8; 16] = [ const PRF_DOMAIN_SEP_DOMAIN_ELEMENT: [u8; 1] = [0x00]; const PRF_DOMAIN_SEP_RANDOMNESS: [u8; 1] = [0x01]; +fn field_from_u128(value: u128) -> F { + // Reduce two 64-bit halves instead of calling the generic 128-bit divider. + const TWO_TO_64_MOD_P: u64 = ((1u128 << 64) % F::ORDER_U64 as u128) as u64; + F::from_u64((value >> 64) as u64) * F::from_u64(TWO_TO_64_MOD_P) + F::from_u64(value as u64) +} + /// A pseudorandom function mapping to field elements. /// It is implemented using Shake128. /// It outputs DOMAIN_LENGTH_FE or RAND_LENGTH_FE many field elements. @@ -68,7 +74,7 @@ where xof_reader.read(&mut buf); // Mapping bytes to a field element - F::from_u128(u128::from_be_bytes(buf)) + field_from_u128(u128::from_be_bytes(buf)) }) } @@ -112,9 +118,32 @@ where xof_reader.read(&mut buf); // Mapping bytes to a field element - F::from_u128(u128::from_be_bytes(buf)) + field_from_u128(u128::from_be_bytes(buf)) }) } + + fn randomness_generator<'a>( + key: &'a Self::Key, + epoch: u32, + message: &'a [u8; crate::MESSAGE_LENGTH], + ) -> impl Fn(u64) -> Self::Randomness + Sync + 'a { + let mut prefix = Shake128::default(); + prefix.update(&PRF_DOMAIN_SEP); + prefix.update(&PRF_DOMAIN_SEP_RANDOMNESS); + prefix.update(key); + prefix.update(&epoch.to_be_bytes()); + prefix.update(message); + move |counter| { + let mut hasher = prefix.clone(); + hasher.update(&counter.to_be_bytes()); + let mut reader = hasher.finalize_xof(); + std::array::from_fn(|_| { + let mut bytes = [0u8; PRF_BYTES_PER_FE]; + reader.read(&mut bytes); + field_from_u128(u128::from_be_bytes(bytes)) + }) + } + } } #[cfg(test)] @@ -127,6 +156,20 @@ mod tests { const RAND_LEN: usize = 4; type PRF = ShakePRFtoF; + proptest! { + #[test] + fn split_reduction_matches_u128(value in any::()) { + prop_assert_eq!(field_from_u128(value), F::from_u128(value)); + } + + #[test] + fn prepared_randomness_matches_scalar(key in prop::array::uniform32(any::()), + message in prop::array::uniform32(any::()), epoch in any::(), counter in any::()) { + let generate = PRF::randomness_generator(&key, epoch, &message); + prop_assert_eq!(generate(counter), PRF::get_randomness(&key, epoch, &message, counter)); + } + } + #[test] fn test_shake_to_field_prf_key_not_all_same() { const K: usize = 10; diff --git a/src/symmetric/tweak_hash.rs b/src/symmetric/tweak_hash.rs index 952b1f7..5d60bb6 100644 --- a/src/symmetric/tweak_hash.rs +++ b/src/symmetric/tweak_hash.rs @@ -48,6 +48,28 @@ pub trait TweakableHash { message: &[Self::Domain], ) -> Self::Domain; + /// Derive chain starts and walk to each signing position. Hash-specific + /// implementations may batch independent chains without changing outputs. + fn sign_chains( + prf_key: &PRF::Key, + parameter: &Self::Parameter, + epoch: u32, + steps: &[u8], + ) -> Vec + where + PRF::Domain: Into, + Self: Sized, + { + steps + .par_iter() + .enumerate() + .map(|(index, &steps)| { + let start = PRF::get_domain_element(prf_key, epoch, index as u64).into(); + chain::(parameter, epoch, index as u8, 0, steps as usize, &start) + }) + .collect() + } + /// Computes one layer of a Merkle tree by hashing pairs of children into parents. /// /// Consecutive pairs of child nodes produce their parent node by hashing diff --git a/src/symmetric/tweak_hash/poseidon.rs b/src/symmetric/tweak_hash/poseidon.rs index 1c43069..4014fc9 100644 --- a/src/symmetric/tweak_hash/poseidon.rs +++ b/src/symmetric/tweak_hash/poseidon.rs @@ -47,26 +47,26 @@ impl PoseidonTweak { level, pos_in_level, } => { - ((*level as u128) << 40) - | ((*pos_in_level as u128) << 8) - | (TWEAK_SEPARATOR_FOR_TREE_HASH as u128) + ((*level as u64) << 40) + | ((*pos_in_level as u64) << 8) + | (TWEAK_SEPARATOR_FOR_TREE_HASH as u64) } Self::ChainTweak { epoch, chain_index, pos_in_chain, } => { - ((*epoch as u128) << 24) - | ((*chain_index as u128) << 16) - | ((*pos_in_chain as u128) << 8) - | (TWEAK_SEPARATOR_FOR_CHAIN_HASH as u128) + ((*epoch as u64) << 24) + | ((*chain_index as u64) << 16) + | ((*pos_in_chain as u64) << 8) + | (TWEAK_SEPARATOR_FOR_CHAIN_HASH as u64) } }; // Now we interpret this integer in base-p to get field elements std::array::from_fn(|_| { - let digit = (acc % F::ORDER_U64 as u128) as u64; - acc /= F::ORDER_U64 as u128; + let digit = acc % F::ORDER_U64; + acc /= F::ORDER_U64; F::from_u64(digit) }) } @@ -270,6 +270,74 @@ pub struct PoseidonTweakHash< const NUM_CHUNKS: usize, >; +impl< + const PARAMETER_LEN: usize, + const HASH_LEN: usize, + const TWEAK_LEN: usize, + const CAPACITY: usize, + const NUM_CHUNKS: usize, +> PoseidonTweakHash +{ + // Compile-time parameter validation for PoseidonTweakHash + // + // DKKW25: https://eprint.iacr.org/2025/055 + const PARAMETER_CHECK: () = { + // The hash output space H = F_p^η and public parameter space + // P = F_p^l_p must both be non-empty. + assert!( + HASH_LEN >= 1, + "Poseidon Tweak Hash: HASH_LEN (η) must be at least 1" + ); + assert!( + PARAMETER_LEN >= 1, + "Poseidon Tweak Hash: PARAMETER_LEN (l_p) must be at least 1" + ); + + // Leaf hashing uses a sponge construction with rate = 24 - CAPACITY. + // The capacity must leave room for at least one rate element. + assert!( + CAPACITY < 24, + "Poseidon Tweak Hash: Capacity must be less than 24" + ); + + // Chain hashing: Th(P, T, M) for M ∈ H uses compression mode + // with width 16. Input = [P || T || M]. + assert!( + PARAMETER_LEN + TWEAK_LEN + HASH_LEN <= 16, + "Poseidon Tweak Chain Hash: Input lengths too large for Poseidon instance" + ); + + // Tree hashing: Th(P, T, M) for M ∈ H² uses compression mode + // with width 24. Input = [P || T || left || right]. + assert!( + PARAMETER_LEN + TWEAK_LEN + 2 * HASH_LEN <= 24, + "Poseidon Tweak Tree Hash: Input lengths too large for Poseidon instance" + ); + + // Leaf hashing: the sponge capacity value V_c is derived from + // 4 domain parameters (each 32 bits = 128 bits total) via a + // width-24 Poseidon compression. + let bits_per_fe = F::ORDER_U64.ilog2() as usize; + assert!( + bits_per_fe * 24 >= DOMAIN_PARAMETERS_LENGTH * 32, + "Poseidon Tweak Leaf Hash: not enough field elements to hash the domain separator" + ); + + // Tweak encoding: the tree tweak packs 32 (pos_in_level) + 8 + // (level) = 40 bits; the chain tweak packs 32 (epoch) + 8 + // (chain_index) + 8 (pos_in_chain) + 8 (separator) = 56 bits. + let tweak_fe_bits = bits_per_fe * TWEAK_LEN; + assert!( + tweak_fe_bits >= 40, + "Poseidon Tweak Hash: not enough field elements to encode the tree tweak" + ); + assert!( + tweak_fe_bits >= 56, + "Poseidon Tweak Hash: not enough field elements to encode the chain tweak" + ); + }; +} + impl< const PARAMETER_LEN: usize, const HASH_LEN: usize, @@ -307,69 +375,84 @@ impl< } } + fn sign_chains( + prf_key: &PRF::Key, + parameter: &Self::Parameter, + epoch: u32, + steps: &[u8], + ) -> Vec + where + PRF::Domain: Into, + { + let () = Self::PARAMETER_CHECK; + assert!(steps.len() <= 256); + assert!(HASH_LEN + PARAMETER_LEN + TWEAK_LEN <= CHAIN_COMPRESSION_WIDTH); + // Group similar chain lengths to reduce inactive SIMD lanes. The + // original indices remain in tweaks and in the returned signature. + let mut order: Vec<_> = (0..steps.len()).collect(); + order.sort_unstable_by_key(|&index| steps[index]); + let groups: Vec<_> = order + .par_chunks(PackedF::WIDTH) + .map(|indices| { + let starts: [Self::Domain; PackedF::WIDTH] = array::from_fn(|lane| { + PRF::get_domain_element(prf_key, epoch, indices[lane % indices.len()] as u64) + .into() + }); + let mut input = [PackedF::ZERO; CHAIN_COMPRESSION_WIDTH]; + input[..HASH_LEN].copy_from_slice(&pack_array(&starts)); + for i in 0..PARAMETER_LEN { + input[HASH_LEN + i] = PackedF::from(parameter[i]); + } + let max_steps = indices.iter().map(|&index| steps[index]).max().unwrap(); + for position in 1..=max_steps { + let tweaks: [[F; TWEAK_LEN]; PackedF::WIDTH] = array::from_fn(|lane| { + Self::chain_tweak(epoch, indices[lane % indices.len()] as u8, position) + .to_field_elements() + }); + for i in 0..TWEAK_LEN { + input[HASH_LEN + PARAMETER_LEN + i] = + PackedF::from_fn(|lane| tweaks[lane][i]); + } + let output = poseidon_compress::( + poseidon1_16(), + &input, + ); + for i in 0..HASH_LEN { + let previous = input[i]; + input[i] = PackedF::from_fn(|lane| { + if position <= steps[indices[lane % indices.len()]] { + output[i].as_slice()[lane] + } else { + previous.as_slice()[lane] + } + }); + } + } + indices + .iter() + .enumerate() + .map(|(lane, &index)| { + ( + index, + FieldArray(array::from_fn(|i| input[i].as_slice()[lane])), + ) + }) + .collect::>() + }) + .collect(); + let mut hashes = vec![FieldArray([F::ZERO; HASH_LEN]); steps.len()]; + for (index, hash) in groups.into_iter().flatten() { + hashes[index] = hash; + } + hashes + } + fn apply( parameter: &Self::Parameter, tweak: &Self::Tweak, message: &[Self::Domain], ) -> Self::Domain { - // Compile-time parameter validation for PoseidonTweakHash - // - // DKKW25: https://eprint.iacr.org/2025/055 - const { - // The hash output space H = F_p^η and public parameter space - // P = F_p^l_p must both be non-empty. - assert!( - HASH_LEN >= 1, - "Poseidon Tweak Hash: HASH_LEN (η) must be at least 1" - ); - assert!( - PARAMETER_LEN >= 1, - "Poseidon Tweak Hash: PARAMETER_LEN (l_p) must be at least 1" - ); - - // Leaf hashing uses a sponge construction with rate = 24 - CAPACITY. - // The capacity must leave room for at least one rate element. - assert!( - CAPACITY < 24, - "Poseidon Tweak Hash: Capacity must be less than 24" - ); - - // Chain hashing: Th(P, T, M) for M ∈ H uses compression mode - // with width 16. Input = [P || T || M]. - assert!( - PARAMETER_LEN + TWEAK_LEN + HASH_LEN <= 16, - "Poseidon Tweak Chain Hash: Input lengths too large for Poseidon instance" - ); - - // Tree hashing: Th(P, T, M) for M ∈ H² uses compression mode - // with width 24. Input = [P || T || left || right]. - assert!( - PARAMETER_LEN + TWEAK_LEN + 2 * HASH_LEN <= 24, - "Poseidon Tweak Tree Hash: Input lengths too large for Poseidon instance" - ); - - // Leaf hashing: the sponge capacity value V_c is derived from - // 4 domain parameters (each 32 bits = 128 bits total) via a - // width-24 Poseidon compression. - let bits_per_fe = F::ORDER_U64.ilog2() as usize; - assert!( - bits_per_fe * 24 >= DOMAIN_PARAMETERS_LENGTH * 32, - "Poseidon Tweak Leaf Hash: not enough field elements to hash the domain separator" - ); - - // Tweak encoding: the tree tweak packs 32 (pos_in_level) + 8 - // (level) = 40 bits; the chain tweak packs 32 (epoch) + 8 - // (chain_index) + 8 (pos_in_chain) + 8 (separator) = 56 bits. - let tweak_fe_bits = bits_per_fe * TWEAK_LEN; - assert!( - tweak_fe_bits >= 40, - "Poseidon Tweak Hash: not enough field elements to encode the tree tweak" - ); - assert!( - tweak_fe_bits >= 56, - "Poseidon Tweak Hash: not enough field elements to encode the chain tweak" - ); - } + let () = Self::PARAMETER_CHECK; // we are in one of three cases: // (1) hashing within chains. We use compression mode. @@ -393,7 +476,7 @@ impl< FieldArray( poseidon_compress::<_, _, CHAIN_COMPRESSION_WIDTH, HASH_LEN>( - &perm, + perm, &combined_input, ), ) @@ -415,7 +498,7 @@ impl< FieldArray( poseidon_compress::<_, _, MERGE_COMPRESSION_WIDTH, HASH_LEN>( - &perm, + perm, &combined_input, ), ) @@ -437,13 +520,13 @@ impl< NUM_CHUNKS as u32, HASH_LEN as u32, ]; - let capacity_value = poseidon_safe_domain_separator::(&perm, &lengths); + let capacity_value = poseidon_safe_domain_separator::(perm, &lengths); FieldArray(poseidon_replacement_sponge::< _, _, MERGE_COMPRESSION_WIDTH, HASH_LEN, - >(&perm, &capacity_value, &combined_input)) + >(perm, &capacity_value, &combined_input)) } _ => FieldArray([F::ONE; HASH_LEN]), // Unreachable case, added for safety } @@ -506,7 +589,7 @@ impl< // Compress all WIDTH parent pairs simultaneously let packed_parents = poseidon_compress::( - &perm, + perm, &packed_input, ); @@ -583,7 +666,7 @@ impl< HASH_LEN as u32, ]; let capacity_val: [PackedF; CAPACITY] = - poseidon_safe_domain_separator::(&sponge_perm, &lengths).map(PackedF::from); + poseidon_safe_domain_separator::(sponge_perm, &lengths).map(PackedF::from); // Compute sponge input length. Required to init packed input vector for each rayon worker let sponge_tweak_offset = PARAMETER_LEN; @@ -673,7 +756,7 @@ impl< // This single call processes all epochs in parallel. *packed_chain = poseidon_compress::( - &chain_perm, + chain_perm, &packed_input, ); } @@ -721,7 +804,7 @@ impl< // Apply the sponge hash to produce the leaf. // This absorbs all chain ends and squeezes out the final hash. poseidon_replacement_sponge::( - &sponge_perm, + sponge_perm, &capacity_val, packed_leaf_input, ) @@ -785,6 +868,36 @@ mod tests { use rand::RngExt; use super::*; + + #[test] + fn signing_chains_match_scalar_for_variable_lengths_and_partial_groups() { + use crate::symmetric::prf::shake_to_field::ShakePRFtoF; + use rand::{SeedableRng, rngs::StdRng}; + type TH = PoseidonTweakHash<5, 7, 2, 9, 20>; + type PRF = ShakePRFtoF<7, 6>; + let mut rng = StdRng::seed_from_u64(43); + let parameter = TH::rand_parameter(&mut rng); + let key = PRF::key_gen(&mut rng); + for count in [0, 1, 3, 4, 5, 20, 39, 78, 155, 256] { + let mut steps: Vec = (0..count).map(|_| rng.random()).collect(); + if count >= 3 { + steps[..3].copy_from_slice(&[0, 1, 255]); + } + let epoch = u32::MAX; + let expected: Vec<_> = steps + .iter() + .enumerate() + .map(|(index, &steps)| { + let start = PRF::get_domain_element(&key, epoch, index as u64).into(); + chain::(¶meter, epoch, index as u8, 0, steps as usize, &start) + }) + .collect(); + assert_eq!( + TH::sign_chains::(&key, ¶meter, epoch, &steps), + expected + ); + } + } use crate::symmetric::prf::shake_to_field::ShakePRFtoF; use p3_field::PrimeField32; use proptest::prelude::*;