Skip to content

refactor(mpt): give digest node references a fixed-size type #680

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Qumeric wants to merge 1 commit into
base: valery/mpt-put-digest
Choose a base branch
from
Draft

refactor(mpt): give digest node references a fixed-size type #680

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 25 additions & 11 deletions crates/mpt/src/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<BranchChildId>`) fits in
/// 4 bytes — `Option<u32>` 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.
Expand Down Expand Up @@ -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
Expand All @@ -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<'_> {
Expand All @@ -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)
}
}
10 changes: 7 additions & 3 deletions crates/mpt/src/resolver.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
};
Expand Down Expand Up @@ -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));
Expand Down
4 changes: 2 additions & 2 deletions crates/mpt/src/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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
Expand Down
91 changes: 34 additions & 57 deletions crates/mpt/src/trie.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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<B: alloy_rlp::BufMut>(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.
Expand Down Expand Up @@ -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)))
}
};

Expand All @@ -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))));
Expand All @@ -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
Expand All @@ -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));
}
Expand Down Expand Up @@ -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()))
}
}
}
Expand Down Expand Up @@ -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);
}
}
}
Expand All @@ -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);
}
}
}
Expand Down Expand Up @@ -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)),
}
}
Expand Down Expand Up @@ -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)
}
Expand Down Expand Up @@ -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))),
}
}

Expand Down Expand Up @@ -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)));
}
};

Expand Down Expand Up @@ -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!(),
};
Expand Down Expand Up @@ -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)));
}
};

Expand All @@ -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));
}
Expand Down Expand Up @@ -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));
}
}
}
Expand Down
Loading