diff --git a/crates/mpt/src/node.rs b/crates/mpt/src/node.rs index a1d9bf259..4883a28bf 100644 --- a/crates/mpt/src/node.rs +++ b/crates/mpt/src/node.rs @@ -4,6 +4,16 @@ use revm_primitives::hex; pub(crate) type NodeId = u32; +/// Length of a node reference that is a keccak digest rather than the node's own encoding. +pub(crate) const DIGEST_LEN: usize = 32; + +/// A node reference that is a keccak digest. +/// +/// The length is part of the type so that comparisons and copies of a digest compile to whole-word +/// operations. A slice of unknown length turns both into `memcmp`/`memcpy` calls, which cost more +/// than the 32 bytes they touch, and the decoder performs one of each per node. +pub(crate) type Digest = [u8; DIGEST_LEN]; + /// Id of a branch child node. `NonZero` so that a child slot (`Option`) fits in /// 4 bytes — `Option` has no niche and takes 8. Node id 0 is the null-node sentinel /// (`NULL_NODE_ID`) and is never stored as a branch child. @@ -33,9 +43,9 @@ pub(crate) enum NodeData<'a> { /// Extension node containing a compact hex-prefix path and a single child. Path encodes a /// shared prefix to skip before continuing at `child`. Extension(&'a [u8], NodeId), - /// Unresolved reference to a node by its Keccak-256 digest (32 bytes). Encountering this in + /// Unresolved reference to a node by its Keccak-256 digest. Encountering this in /// `get`/`insert`/`delete` is an error; resolution happens in `build_mpt` helpers. - Digest(&'a [u8]), + Digest(&'a Digest), } /// Represents the ways in which one node can reference another node inside the sparse Merkle @@ -49,9 +59,8 @@ pub(crate) enum NodeRef<'a> { /// used for short encodings that are less than 32 bytes in length. Bytes(&'a [u8]), /// Represents an indirect reference to another node using the Keccak hash of its long - /// encoding, so its length is always 32. Used for encodings that are not less than 32 bytes in - /// length. - Digest(&'a [u8]), + /// encoding. Used for encodings that are not less than [`DIGEST_LEN`] bytes in length. + Digest(&'a Digest), } impl core::fmt::Display for NodeRef<'_> { @@ -75,17 +84,22 @@ impl<'a> NodeRef<'a> { pub(crate) fn as_slice(&self) -> &'a [u8] { match self { NodeRef::Bytes(slice) => slice, - NodeRef::Digest(slice) => slice, + NodeRef::Digest(digest) => digest.as_slice(), } } + /// Interprets one RLP-encoded reference: either a digest behind its one-byte string header, or + /// a short node's own encoding. + /// + /// The length is tested first so that a short reference — the common case, and one per branch + /// child — costs a single comparison. #[inline(always)] pub(crate) fn from_rlp_slice(slice: &'a [u8]) -> Self { - if slice.len() == 33 { - Self::Digest(&slice[1..]) - } else { - debug_assert!(slice.len() < 32); - Self::Bytes(slice) + if slice.len() == DIGEST_LEN + 1 { + if let Ok(digest) = <&Digest>::try_from(&slice[1..]) { + return Self::Digest(digest); + } } + Self::Bytes(slice) } } diff --git a/crates/mpt/src/resolver.rs b/crates/mpt/src/resolver.rs index add478b02..daf170857 100644 --- a/crates/mpt/src/resolver.rs +++ b/crates/mpt/src/resolver.rs @@ -1,5 +1,5 @@ use crate::{ - node::{BranchChildId, NodeData, NodeId}, + node::{BranchChildId, NodeData, NodeId, DIGEST_LEN}, trie::{NULL_NODE_ID, NULL_NODE_REF_SLICE}, Error, Mpt, }; @@ -48,11 +48,15 @@ impl MptResolver { let node_id = match alloy_rlp::Header::decode_raw(node_bytes)? { PayloadView::String(item) => match item.len() { 0 => NULL_NODE_ID, - 32 => match self.node_store.get(&B256::from_slice(item)) { + DIGEST_LEN => match self.node_store.get(&B256::from_slice(item)) { Some(resolved_node_bytes) => { self.resolve_internal(&mut resolved_node_bytes.as_ref(), mpt)? } - None => mpt.add_node_copied(&NodeData::Digest(item)), + None => match item.try_into() { + Ok(digest) => mpt.add_node_copied(&NodeData::Digest(digest)), + // Unreachable: the arm matched on the length. + Err(_) => return Err(Error::RlpError(alloy_rlp::Error::UnexpectedLength)), + }, }, _ => { return Err(Error::RlpError(alloy_rlp::Error::UnexpectedLength)); diff --git a/crates/mpt/src/tests.rs b/crates/mpt/src/tests.rs index b4dc3a8fc..17f1b1f45 100644 --- a/crates/mpt/src/tests.rs +++ b/crates/mpt/src/tests.rs @@ -277,7 +277,7 @@ fn test_serde_rejects_corrupted_digest_reference() -> Result<(), Error> { #[test] fn test_serde_digest_root() -> Result<(), Error> { let bump = bumpalo::Bump::new(); - let digest = bump.alloc_slice_copy(&[0xabu8; 32]); + let digest = bump.alloc([0xabu8; 32]); let mut trie = Mpt::new(&bump); let root_id = trie.add_node(NodeData::Digest(digest), None); trie.set_root_id(root_id); @@ -301,7 +301,7 @@ fn test_delete_with_unresolved_sibling_errors() { let mut trie = Mpt::new(&bump); // Create a fake 32-byte digest (simulating hash of an unknown node) - let fake_digest: &[u8] = bump.alloc_slice_copy(&[0xABu8; 32]); + let fake_digest: &[u8; 32] = bump.alloc([0xABu8; 32]); // Build structure: Branch -> [Leaf at index 0, Digest at index 1] // When we delete the Leaf, the Branch should collapse, but we don't know diff --git a/crates/mpt/src/trie.rs b/crates/mpt/src/trie.rs index 07a2c8c66..d193394bd 100644 --- a/crates/mpt/src/trie.rs +++ b/crates/mpt/src/trie.rs @@ -14,15 +14,12 @@ use crate::{ encoded_path_eq_key, encoded_path_from_key, encoded_path_strip_prefix_key, lcp_key, prefix_to_nibs, to_encoded_path_with_bump, KeyNibbles, }, - node::{BranchChildId, BranchChildren, NodeData, NodeId, NodeRef}, + node::{BranchChildId, BranchChildren, Digest, NodeData, NodeId, NodeRef, DIGEST_LEN}, }; /// OpenVM memory alignment word size. const MIN_ALIGN: usize = 4; -/// Length of a node reference that is a keccak digest rather than the node's own encoding. -const DIGEST_LEN: usize = 32; - /// Sentinel index representing the null node when decoding and in internal references. /// In a default MPT, `nodes[0]` starts as `Null`, but the root may later be changed to a /// non-null node (e.g. `Digest`) for convenience. `NULL_NODE_ID` is still used by the decoder @@ -152,41 +149,18 @@ fn is_null_ref(slice: &[u8]) -> bool { } /// Byte-slice equality as an explicit loop. Slice `==` on slices of unknown length compiles to a -/// `memcmp` call; the node references compared during decoding are at most 33 bytes, where the -/// call overhead dominates an inline loop. Prefer [`digest_eq`] whenever one side is a full-length -/// digest, since a length known at compile time turns the comparison into whole-word loads. +/// `memcmp` call, and the short node references compared during decoding are small enough that the +/// call dominates an inline loop. Digests do not need this: their length is part of [`Digest`], so +/// comparing them with `==` stays inline as whole-word loads. #[inline(always)] fn bytes_eq(a: &[u8], b: &[u8]) -> bool { a.len() == b.len() && core::iter::zip(a, b).all(|(x, y)| x == y) } -/// Equality against a 32-byte node reference. -/// -/// Fixing the length lets the comparison stay inline as four word loads per side and an XOR -/// chain, where a loop over slices of unknown length pays several instructions for every byte. -/// Slices of different lengths are never equal, so a reference that is not a full digest is -/// inequality rather than a case to fall back on. +/// Reads a digest from `slice`, which must be exactly [`DIGEST_LEN`] bytes. #[inline(always)] -fn digest_eq(a: &[u8], b: &[u8]) -> bool { - match (<&[u8; DIGEST_LEN]>::try_from(a), <&[u8; DIGEST_LEN]>::try_from(b)) { - (Ok(a), Ok(b)) => a == b, - _ => false, - } -} - -/// Appends a 32-byte node reference to `out`. -/// -/// Fixing the length keeps the copy inline as whole-word moves. Handing the buffer a slice whose -/// length the compiler cannot see turns the copy into a `memcpy` call, and at this size the call -/// costs more than the bytes it moves — encoding node references is the decoder's busiest copy. -#[inline(always)] -fn put_digest(digest: &[u8], out: &mut B) { - match <&[u8; DIGEST_LEN]>::try_from(digest) { - Ok(digest) => out.put_slice(digest), - // A digest reference always holds `DIGEST_LEN` bytes, so this arm exists only to keep the - // function total without a conversion that could panic. - Err(_) => out.put_slice(digest), - } +fn digest_from(slice: &[u8]) -> Result<&Digest, Error> { + slice.try_into().map_err(|_| Error::RlpError(alloy_rlp::Error::UnexpectedLength)) } /// Converts a node id to a branch child slot id. @@ -303,11 +277,10 @@ impl<'a> Mpt<'a> { if rlp_node.len() < 32 { NodeRef::Bytes(rlp_node) } else if !list { - NodeRef::Digest(payload) + NodeRef::Digest(digest_from(payload)?) } else { - let digest = keccak256(rlp_node); - let digest_slice = bump.alloc_slice_copy(digest.as_slice()); - NodeRef::Digest(digest_slice) + // Allocating the array rather than a slice of it keeps the copy a fixed-size move. + NodeRef::Digest(bump.alloc(keccak256(rlp_node))) } }; @@ -330,8 +303,8 @@ impl<'a> Mpt<'a> { // the declared RLP payload and its alignment padding. let encoded = unsafe { advance_unchecked(bytes, 33) }; unsafe { advance_unchecked(bytes, 3) }; - let digest = &encoded[1..]; - if !digest_eq(digest, expected_node_ref.as_slice()) { + let digest = digest_from(&encoded[1..])?; + if digest.as_slice() != expected_node_ref.as_slice() { return Err(Error::NodeRefMismatch); } return Ok(self.add_node(NodeData::Digest(digest), Some(NodeRef::Digest(digest)))); @@ -358,13 +331,13 @@ impl<'a> Mpt<'a> { } NodeRef::Bytes(rlp_node) } else if payload_length == 32 && !list { - if !digest_eq(payload, expected_node_ref.as_slice()) { + if payload != expected_node_ref.as_slice() { return Err(Error::NodeRefMismatch); } expected_node_ref } else { let digest = keccak256(rlp_node); - if !digest_eq(digest.as_slice(), expected_node_ref.as_slice()) { + if digest.as_slice() != expected_node_ref.as_slice() { return Err(Error::NodeRefMismatch); } expected_node_ref @@ -374,7 +347,10 @@ impl<'a> Mpt<'a> { if !list { let node_id = match payload_length { 0 => NULL_NODE_ID, - 32 => self.add_node(NodeData::Digest(payload), Some(NodeRef::Digest(payload))), + DIGEST_LEN => { + let digest = digest_from(payload)?; + self.add_node(NodeData::Digest(digest), Some(NodeRef::Digest(digest))) + } _ => { return Err(Error::RlpError(alloy_rlp::Error::UnexpectedLength)); } @@ -495,9 +471,7 @@ impl<'a> Mpt<'a> { self.encode_with_payload_len(node_id, payload_length, &mut sponge); debug_assert_eq!(sponge.absorbed_len(), rlp_length); - let digest = sponge.finalize(); - let digest_slice = self.bump.alloc_slice_copy(&digest); - NodeRef::Digest(digest_slice) + NodeRef::Digest(self.bump.alloc(sponge.finalize())) } } } @@ -536,7 +510,7 @@ impl<'a> Mpt<'a> { self.reference_encode(*child_id, out); } NodeData::Digest(digest) => { - encode_slice(digest, out); + encode_slice(digest.as_slice(), out); } } } @@ -557,8 +531,8 @@ impl<'a> Mpt<'a> { NodeRef::Bytes(bytes) => out.put_slice(bytes), // if the reference is a digest, RLP-encode it with its fixed known length NodeRef::Digest(digest) => { - out.put_u8(alloy_rlp::EMPTY_STRING_CODE + 32); - put_digest(digest, out); + out.put_u8(alloy_rlp::EMPTY_STRING_CODE + DIGEST_LEN as u8); + out.put_slice(digest); } } } @@ -619,7 +593,7 @@ impl<'a> Mpt<'a> { } }; match node_ref { - NodeRef::Digest(digest) => B256::from_slice(digest), + NodeRef::Digest(digest) => B256::new(*digest), NodeRef::Bytes(bytes) => B256::new(keccak256(bytes)), } } @@ -725,7 +699,7 @@ impl<'a> Mpt<'a> { NodeData::Extension(prefix, ext_node_id) => { NodeData::Extension(self.bump.alloc_slice_copy(prefix), *ext_node_id) } - NodeData::Digest(digest) => NodeData::Digest(self.bump.alloc_slice_copy(digest)), + NodeData::Digest(digest) => NodeData::Digest(self.bump.alloc(**digest)), }; self.add_node(data, None) } @@ -769,7 +743,7 @@ impl<'a> Mpt<'a> { Ok(None) } } - NodeData::Digest(digest) => Err(Error::NodeNotResolved(B256::from_slice(digest))), + NodeData::Digest(digest) => Err(Error::NodeNotResolved(B256::new(**digest))), } } @@ -891,7 +865,7 @@ impl<'a> Mpt<'a> { } } NodeData::Digest(digest) => { - return Err(Error::NodeNotResolved(B256::from_slice(digest))); + return Err(Error::NodeNotResolved(B256::new(*digest))); } }; @@ -963,7 +937,7 @@ impl<'a> Mpt<'a> { NodeData::Extension(new_path, child_id) } NodeData::Digest(digest) => { - return Err(Error::NodeNotResolved(B256::from_slice(digest))); + return Err(Error::NodeNotResolved(B256::new(*digest))); } NodeData::Null => unreachable!(), }; @@ -1020,14 +994,14 @@ impl<'a> Mpt<'a> { NodeData::Branch(_) => NodeData::Extension(prefix, child_id), // for a digest, we don't know the node type so we can't safely canonicalize NodeData::Digest(digest) => { - return Err(Error::NodeNotResolved(B256::from_slice(digest))); + return Err(Error::NodeNotResolved(B256::new(**digest))); } }; self.nodes[node_id as usize] = new_node_data; true } NodeData::Digest(digest) => { - return Err(Error::NodeNotResolved(B256::from_slice(digest))); + return Err(Error::NodeNotResolved(B256::new(*digest))); } }; @@ -1052,7 +1026,10 @@ impl<'a> Mpt<'a> { let node_id = match alloy_rlp::Header::decode_raw(bytes)? { alloy_rlp::PayloadView::String(item) => match item.len() { 0 => NULL_NODE_ID, - 32 => self.add_node(NodeData::Digest(item), Some(NodeRef::Digest(item))), + DIGEST_LEN => { + let digest = digest_from(item)?; + self.add_node(NodeData::Digest(digest), Some(NodeRef::Digest(digest))) + } _ => { return Err(Error::RlpError(alloy_rlp::Error::UnexpectedLength)); } @@ -1155,7 +1132,7 @@ impl Mpt<'_> { self.print_trie_internal(*child_id, depth + 1); } NodeData::Digest(digest) => { - println!("{}Digest {:?}", indent, B256::from_slice(digest)); + println!("{}Digest {:?}", indent, B256::new(**digest)); } } }