diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5258ce7..b91cbae 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -80,6 +80,8 @@ jobs: args: "" - features: alloc-only args: "--no-default-features --features alloc" + - features: hashtree + args: "--features libssz-merkle/hashtree" steps: - uses: actions/checkout@v7 - uses: actions/cache@v6 diff --git a/crates/ssz-merkle/Cargo.toml b/crates/ssz-merkle/Cargo.toml index cae1346..4c83e59 100644 --- a/crates/ssz-merkle/Cargo.toml +++ b/crates/ssz-merkle/Cargo.toml @@ -14,12 +14,16 @@ default = ["std"] std = ["alloc", "libssz/std", "sha2-backend"] alloc = ["libssz/alloc"] sha2-backend = ["dep:sha2"] +# Hash Merkle pairs with the `hashtree` library: multi-buffer SHA-256 assembly +# specialized for 64-byte inputs. Needs a C toolchain and `make` at build time. +hashtree = ["std", "sha2-backend", "dep:hashtree-rs"] ethereum_types = ["dep:ethereum-types", "libssz/ethereum_types"] [dependencies] libssz = { version = "0.3.0", path = "../ssz", default-features = false } -sha2 = { version = "0.10", default-features = false, optional = true } +sha2 = { version = "0.11", default-features = false, optional = true } +hashtree-rs = { version = "0.2.6", optional = true } ethereum-types = { version = "0.15", default-features = false, optional = true } [build-dependencies] -sha2 = "0.10" +sha2 = "0.11" diff --git a/crates/ssz-merkle/src/lib.rs b/crates/ssz-merkle/src/lib.rs index f175298..558bb42 100644 --- a/crates/ssz-merkle/src/lib.rs +++ b/crates/ssz-merkle/src/lib.rs @@ -16,33 +16,43 @@ pub type Node = [u8; 32]; include!(concat!(env!("OUT_DIR"), "/zero_hashes.rs")); +#[cfg(feature = "sha2-backend")] +mod sha256; +#[cfg(feature = "sha2-backend")] +pub use sha256::Sha2Hasher; + /// Trait for pluggable SHA-256 implementations. /// /// Instance method (`&self`) enables passing `&impl Sha256Hasher` through the /// call chain. The hasher struct is typically zero-sized so `&self` compiles away. pub trait Sha256Hasher { + /// SHA-256 of arbitrary bytes. fn hash(&self, data: &[u8]) -> [u8; 32]; -} -/// Default SHA-256 hasher using the `sha2` crate. -#[cfg(feature = "sha2-backend")] -pub struct Sha2Hasher; - -#[cfg(feature = "sha2-backend")] -impl Sha256Hasher for Sha2Hasher { - fn hash(&self, data: &[u8]) -> [u8; 32] { - use sha2::{Digest, Sha256}; - Sha256::digest(data).into() + /// Hash each pair, `SHA256(pair[0] || pair[1])`, into the matching slot of `out`. + /// + /// This is the hot path of Merkleization: [`merkleize`] hashes a whole tree + /// level per call and [`hash_nodes`] passes a single pair, so an + /// implementation can batch pairs or specialize for 64-byte messages. The + /// default hashes one pair at a time through [`hash`](Self::hash). + /// + /// # Panics + /// + /// Panics if `pairs.len() != out.len()`. + fn hash_pairs(&self, pairs: &[[Node; 2]], out: &mut [Node]) { + assert_eq!(pairs.len(), out.len(), "one output node per pair"); + for (pair, out) in pairs.iter().zip(out) { + *out = self.hash(pair.as_flattened()); + } } } /// Hash two nodes together: SHA256(a || b). #[inline] pub fn hash_nodes(hasher: &impl Sha256Hasher, a: &Node, b: &Node) -> Node { - let mut buf = [0u8; 64]; - buf[..32].copy_from_slice(a); - buf[32..].copy_from_slice(b); - hasher.hash(&buf) + let mut out = [[0u8; 32]]; + hasher.hash_pairs(&[[*a, *b]], &mut out); + out[0] } /// Pack serialized bytes into 32-byte chunks, zero-padding the last chunk if needed. @@ -79,7 +89,9 @@ pub fn pack_bits(bits: &[u8], num_bits: usize) -> Vec { /// all zeros, so its root is read from [`ZERO_HASHES`] instead of being hashed. /// /// Cost is `O(chunks.len())` plus one hash per virtual level above the real subtree. -/// Two buffers are allocated and reused (≈ `chunks.len()` and `chunks.len().div_ceil(2)` nodes). +/// Each tree level is hashed in one [`Sha256Hasher::hash_pairs`] batch. Two buffers are +/// allocated and reused (≈ `chunks.len() / 2` and `chunks.len() / 4` nodes); the leaves are +/// never copied. #[cfg(feature = "alloc")] pub fn merkleize(hasher: &impl Sha256Hasher, chunks: &[Node], limit: Option) -> Node { let count = match limit { @@ -110,36 +122,46 @@ pub fn merkleize(hasher: &impl Sha256Hasher, chunks: &[Node], limit: Option = chunks.to_vec(); - let mut next: Vec = Vec::with_capacity(layer.len().div_ceil(2)); - - for zero_hash in ZERO_HASHES.iter().take(real_depth) { - next.clear(); - let (pairs, remainder) = layer.as_chunks::<2>(); - for pair in pairs { - next.push(hash_nodes(hasher, &pair[0], &pair[1])); + // Fold the real chunks one level at a time. Level 0 reads `chunks` in place; + // later levels swap between two buffers allocated once, so no level allocates. + let mut root = if real_depth == 0 { + chunks[0] + } else { + let mut layer: Vec = Vec::with_capacity(chunks.len().div_ceil(2)); + hash_level(hasher, chunks, &ZERO_HASHES[0], &mut layer); + let mut next: Vec = Vec::with_capacity(layer.len().div_ceil(2)); + for zero_hash in &ZERO_HASHES[1..real_depth] { + hash_level(hasher, &layer, zero_hash, &mut next); + core::mem::swap(&mut layer, &mut next); } - if let [odd] = remainder { - next.push(hash_nodes(hasher, odd, zero_hash)); - } - core::mem::swap(&mut layer, &mut next); - } + layer[0] + }; - // `layer[0]` now roots the real-data subtree. Merge in the virtual levels. - let mut root = layer[0]; - for zero_hash in ZERO_HASHES.iter().take(depth).skip(real_depth) { + // `root` now roots the real-data subtree. Merge in the virtual levels. + for zero_hash in &ZERO_HASHES[real_depth..depth] { root = hash_nodes(hasher, &root, zero_hash); } root } +/// Hash one tree level of `layer` into `next`. +/// +/// Full pairs of real nodes are hashed in a single [`Sha256Hasher::hash_pairs`] +/// batch. A level with an odd node count ends on the frontier, where the missing +/// right sibling roots an all-zero subtree, so `zero_hash` (the precomputed root +/// of that subtree) stands in for it rather than being hashed. +#[cfg(feature = "alloc")] +fn hash_level(hasher: &impl Sha256Hasher, layer: &[Node], zero_hash: &Node, next: &mut Vec) { + let (pairs, remainder) = layer.as_chunks::<2>(); + next.clear(); + next.resize(pairs.len(), [0u8; 32]); + hasher.hash_pairs(pairs, next); + if let [odd] = remainder { + next.push(hash_nodes(hasher, odd, zero_hash)); + } +} + /// Progressive merkleization: a chain of progressively larger binary subtrees. /// /// Creates a 0-terminated sequence of binary subtrees with increasing leaf count: diff --git a/crates/ssz-merkle/src/sha256.rs b/crates/ssz-merkle/src/sha256.rs new file mode 100644 index 0000000..6e276c0 --- /dev/null +++ b/crates/ssz-merkle/src/sha256.rs @@ -0,0 +1,138 @@ +//! Default SHA-256 backend: the `sha2` crate, or `hashtree` with the `hashtree` feature. + +use crate::{Node, Sha256Hasher}; + +/// Default SHA-256 hasher. +/// +/// Merkle pairs are hashed by calling the SHA-256 compression function directly +/// on the 64-byte message and a constant padding block, skipping the +/// general-purpose `Digest` buffering. `sha2` selects CPU SHA instructions at +/// runtime (x86 SHA-NI, ARMv8 SHA2) and falls back to software otherwise. +/// +/// With the `hashtree` feature, pairs are hashed by the +/// [`hashtree`](https://github.com/OffchainLabs/hashtree) library instead: +/// assembly specialized for 64-byte inputs that also hashes several pairs at +/// once (SHA-NI, AVX-512 x16, AVX2 x8, ARMv8 SHA2, NEON x4). +pub struct Sha2Hasher; + +impl Sha256Hasher for Sha2Hasher { + fn hash(&self, data: &[u8]) -> [u8; 32] { + use sha2::{Digest, Sha256}; + Sha256::digest(data).into() + } + + #[cfg(not(feature = "hashtree"))] + fn hash_pairs(&self, pairs: &[[Node; 2]], out: &mut [Node]) { + assert_eq!(pairs.len(), out.len(), "one output node per pair"); + for (pair, out) in pairs.iter().zip(out) { + *out = hash_pair(pair); + } + } + + #[cfg(feature = "hashtree")] + fn hash_pairs(&self, pairs: &[[Node; 2]], out: &mut [Node]) { + // `hashtree_rs::hash` performs no length checks of its own. + assert_eq!(pairs.len(), out.len(), "one output node per pair"); + if pairs.is_empty() { + return; + } + // The C library selects its implementation lazily by writing a static + // function pointer on first use; do that once, before any concurrent call. + static INIT: std::sync::Once = std::sync::Once::new(); + INIT.call_once(|| { + hashtree_rs::init(); + }); + hashtree_rs::hash( + out.as_flattened_mut(), + pairs.as_flattened().as_flattened(), + pairs.len(), + ); + } +} + +/// SHA-256 initial hash value (FIPS 180-4, 5.3.3). +const IV: [u32; 8] = [ + 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19, +]; + +/// Padding for a 64-byte message, which always fills a block of its own: +/// `0x80`, zeros, then the message length in bits (512) as a big-endian `u64`. +const PADDING_BLOCK: [u8; 64] = { + let mut block = [0u8; 64]; + block[0] = 0x80; + block[62] = 0x02; + block +}; + +/// SHA-256 of `pair[0] || pair[1]`, as two calls to the compression function. +#[cfg_attr(feature = "hashtree", allow(dead_code))] +#[inline] +fn hash_pair(pair: &[Node; 2]) -> Node { + let mut message = [0u8; 64]; + message[..32].copy_from_slice(&pair[0]); + message[32..].copy_from_slice(&pair[1]); + + let mut state = IV; + sha2::block_api::compress256(&mut state, &[message, PADDING_BLOCK]); + + let mut digest = [0u8; 32]; + let (words, _) = digest.as_chunks_mut::<4>(); + for (bytes, word) in words.iter_mut().zip(state) { + *bytes = word.to_be_bytes(); + } + digest +} + +#[cfg(all(test, feature = "alloc"))] +mod tests { + use super::*; + use alloc::vec::Vec; + use sha2::{Digest, Sha256}; + + /// Pairs whose bytes differ in every position and set high bits, so a + /// misplaced word or byte-order slip changes the digest. + fn sample_pairs(count: usize) -> Vec<[Node; 2]> { + (0..count) + .map(|i| { + let mut pair = [[0u8; 32]; 2]; + for (j, byte) in pair.as_flattened_mut().iter_mut().enumerate() { + *byte = (i * 131 + j * 29 + 0x80) as u8; + } + pair + }) + .collect() + } + + fn digest(pair: &[Node; 2]) -> Node { + Sha256::digest(pair.as_flattened()).into() + } + + #[test] + fn hash_pair_matches_digest() { + for pair in sample_pairs(16) { + assert_eq!(hash_pair(&pair), digest(&pair)); + } + } + + /// Covers the batch path used by `merkleize`, including the `hashtree` + /// backend when that feature is enabled, across every SIMD batch remainder. + #[test] + fn hash_pairs_matches_digest_for_every_count() { + for count in 0..=33 { + let pairs = sample_pairs(count); + let mut out = alloc::vec![[0u8; 32]; count]; + Sha2Hasher.hash_pairs(&pairs, &mut out); + for (i, (pair, node)) in pairs.iter().zip(&out).enumerate() { + assert_eq!(*node, digest(pair), "pair {i} of {count}"); + } + } + } + + #[test] + #[should_panic(expected = "one output node per pair")] + fn hash_pairs_rejects_mismatched_output() { + let pairs = sample_pairs(3); + let mut out = [[0u8; 32]; 2]; + Sha2Hasher.hash_pairs(&pairs, &mut out); + } +} diff --git a/docs/technical-design.md b/docs/technical-design.md index d1bc16e..a417b55 100644 --- a/docs/technical-design.md +++ b/docs/technical-design.md @@ -196,6 +196,8 @@ merkleize(chunks, limit): **Virtual padding:** rather than allocating `limit` chunks, the implementation uses precomputed `ZERO_HASHES[level]` for any subtree that would be all zeros. This keeps memory usage proportional to the actual number of non-zero chunks. +**Batched levels:** every full pair of a level is handed to `Sha256Hasher::hash_pairs` in one call, and `hash_nodes` passes a single pair through the same method. A hasher can therefore hash many 64-byte messages at once (multi-buffer SIMD) or skip general-purpose padding. The default `Sha2Hasher` calls the `sha2` compression function directly on the message and a constant padding block (`sha2` picks SHA-NI / ARMv8 SHA2 at runtime). The optional `hashtree` feature hands each batch to the [hashtree](https://github.com/OffchainLabs/hashtree) assembly library instead, at the cost of a C toolchain at build time. + ### Zero hash table `ZERO_HASHES` is a static array of 65 entries (depth 0 through 64), computed at build time by `ssz-merkle/build.rs`: @@ -298,4 +300,4 @@ Errors are data, not strings. Each variant carries the numeric values needed to | Derive macros | Separate crate | Separate crate | Integrated | | Zero hashes | Build-time precomputed | `lazy_static` | Computed at init | | Error types | Flat enum | Flat enum | Flat enum | -| Hash function | sha2 crate | ring / sha2 | sha2 | +| Hash function | sha2 crate (optional hashtree) | ring / sha2 | sha2 |