feat: leaf node decoding

src/hash_builder/mod.rs

@@ -1,7 +1,7 @@
 //! The implementation of the hash builder.
 
 use super::{
-    nodes::{word_rlp, BranchNode, ExtensionNode, LeafNode},
+    nodes::{word_rlp, BranchNode, ExtensionNode, LeafNodeRef},
     BranchNodeCompact, Nibbles, TrieMask, EMPTY_ROOT_HASH,
 };
 use crate::HashMap;
@@ -226,7 +226,7 @@ impl HashBuilder {
             if !build_extensions {
                 match &self.value {
                     HashBuilderValue::Bytes(leaf_value) => {
-                        let leaf_node = LeafNode::new(&short_node_key, leaf_value);
+                        let leaf_node = LeafNodeRef::new(&short_node_key, leaf_value);
                         trace!(target: "trie::hash_builder", ?leaf_node, "pushing leaf node");
                         trace!(target: "trie::hash_builder", rlp = {
                             self.rlp_buf.clear();
@@ -407,6 +407,7 @@ impl HashBuilder {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::nodes::LeafNode;
     use alloy_primitives::{b256, hex, U256};
     use alloy_rlp::Encodable;
 
@@ -584,20 +585,19 @@ mod tests {
             (hex!("646f").to_vec(), hex!("76657262").to_vec()),
             (hex!("676f6f64").to_vec(), hex!("7075707079").to_vec()),
         ];
-        let input =
-            raw_input.iter().map(|(key, value)| (Nibbles::unpack(key), value)).collect::<Vec<_>>();
+        let expected = triehash_trie_root(raw_input.clone());
 
         // We create the hash builder and add the leaves
         let mut hb = HashBuilder::default();
-        for (key, val) in input.iter() {
-            hb.add_leaf(key.clone(), val.as_slice());
+        for (key, val) in &raw_input {
+            hb.add_leaf(Nibbles::unpack(key), val.as_slice());
         }
 
         // Manually create the branch node that should be there after the first 2 leaves are added.
         // Skip the 0th element given in this example they have a common prefix and will
         // collapse to a Branch node.
-        let leaf1 = LeafNode::new(&Nibbles::unpack(&raw_input[0].0[1..]), input[0].1);
-        let leaf2 = LeafNode::new(&Nibbles::unpack(&raw_input[1].0[1..]), input[1].1);
+        let leaf1 = LeafNode::new(Nibbles::unpack(&raw_input[0].0[1..]), raw_input[0].1.clone());
+        let leaf2 = LeafNode::new(Nibbles::unpack(&raw_input[1].0[1..]), raw_input[1].1.clone());
         let mut branch: [&dyn Encodable; 17] = [b""; 17];
         // We set this to `4` and `7` because that mathces the 2nd element of the corresponding
         // leaves. We set this to `7` because the 2nd element of Leaf 1 is `7`.
@@ -611,7 +611,6 @@ mod tests {
         // Insert the branch with the `0x6` shared prefix.
         hb2.add_branch(Nibbles::from_nibbles_unchecked([0x6]), branch_node_hash, false);
 
-        let expected = triehash_trie_root(raw_input.clone());
         assert_eq!(hb.root(), expected);
         assert_eq!(hb2.root(), expected);
     }

src/nodes/leaf.rs

@@ -1,7 +1,7 @@
 use super::{super::Nibbles, rlp_node};
-use alloy_rlp::{BufMut, Encodable};
+use alloy_primitives::Bytes;
+use alloy_rlp::{length_of_length, BufMut, Decodable, Encodable, Header, EMPTY_STRING_CODE};
 use core::fmt;
-use smallvec::SmallVec;
 
 #[allow(unused_imports)]
 use alloc::vec::Vec;
@@ -13,49 +13,122 @@ use alloc::vec::Vec;
 /// remaining portion of the key after following the path through the trie, and the value is the
 /// data associated with the full key. When searching the trie for a specific key, reaching a leaf
 /// node means that the search has successfully found the value associated with that key.
-#[derive(Default)]
-pub struct LeafNode<'a> {
-    /// The key path. See [`Nibbles::encode_path_leaf`] for more information.
-    pub key: SmallVec<[u8; 36]>,
+#[derive(PartialEq, Eq, Debug)]
+pub struct LeafNode {
+    /// The key for this leaf node.
+    pub key: Nibbles,
     /// The node value.
-    pub value: &'a [u8],
+    pub value: Vec<u8>,
 }
 
-impl<'a> LeafNode<'a> {
+impl LeafNode {
     /// Creates a new leaf node with the given key and value.
-    pub fn new(key: &Nibbles, value: &'a [u8]) -> Self {
-        Self { key: key.encode_path_leaf(true), value }
+    pub fn new(key: Nibbles, value: Vec<u8>) -> Self {
+        Self { key, value }
     }
 
-    /// RLP encodes the node and returns either RLP(Node) or RLP(keccak(RLP(node)))
-    /// depending on if the serialized node was longer than a keccak).
-    pub fn rlp(&self, out: &mut Vec<u8>) -> Vec<u8> {
-        self.encode(out);
-        rlp_node(out)
+    /// Return leaf node as [LeafNodeRef].
+    pub fn as_ref(&self) -> LeafNodeRef<'_> {
+        LeafNodeRef { key: &self.key, value: &self.value }
     }
 }
 
-// Handroll because `key` must be encoded as a slice
-impl Encodable for LeafNode<'_> {
+impl Encodable for LeafNode {
     fn encode(&self, out: &mut dyn BufMut) {
-        #[derive(alloy_rlp::RlpEncodable)]
-        struct S<'a> {
-            encoded_path: &'a [u8],
-            value: &'a [u8],
+        self.as_ref().encode(out)
+    }
+}
+
+impl Decodable for LeafNode {
+    fn decode(buf: &mut &[u8]) -> alloy_rlp::Result<Self> {
+        let mut bytes = Header::decode_bytes(buf, true)?;
+
+        let encoded_key = Bytes::decode(&mut bytes)?;
+        if encoded_key.is_empty() {
+            return Err(alloy_rlp::Error::Custom("leaf node key empty"));
         }
-        S { encoded_path: &self.key, value: self.value }.encode(out);
+
+        // Retrieve first byte. If it's [Some], then the nibbles are odd.
+        let first = match encoded_key[0] & 0xf0 {
+            0x30 => Some(encoded_key[0] & 0x0f),
+            0x20 => None,
+            _ => return Err(alloy_rlp::Error::Custom("node is not leaf")),
+        };
+        let rest = encoded_key[1..].iter().flat_map(|b| [b >> 4, b & 0x0f]);
+
+        let is_odd = first.is_some();
+        let len = (encoded_key.len() - 1) * 2 + is_odd as usize;
+        let mut nibbles = Vec::with_capacity(len);
+        unsafe {
+            let ptr: *mut u8 = nibbles.as_mut_ptr();
+            for (i, nibble) in first.into_iter().chain(rest).enumerate() {
+                ptr.add(i).write(nibble)
+            }
+            nibbles.set_len(len);
+        }
+
+        Ok(Self {
+            key: Nibbles::from_vec_unchecked(nibbles),
+            value: Bytes::decode(&mut bytes)?.to_vec(),
+        })
     }
 }
 
-impl fmt::Debug for LeafNode<'_> {
+/// Reference to the leaf node. See [LeafNode] from more information.
+pub struct LeafNodeRef<'a> {
+    /// The key for this leaf node.
+    pub key: &'a Nibbles,
+    /// The node value.
+    pub value: &'a [u8],
+}
+
+impl fmt::Debug for LeafNodeRef<'_> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("LeafNode")
-            .field("key", &alloy_primitives::hex::encode(&self.key))
+            .field("key", &self.key)
             .field("value", &alloy_primitives::hex::encode(self.value))
             .finish()
     }
 }
 
+/// Manual implementation of encoding for the leaf node of Merkle Patricia Trie.
+impl Encodable for LeafNodeRef<'_> {
+    fn encode(&self, out: &mut dyn BufMut) {
+        Header { list: true, payload_length: self.rlp_payload_length() }.encode(out);
+        self.key.encode_path_leaf(true).as_slice().encode(out);
+        self.value.encode(out);
+    }
+
+    fn length(&self) -> usize {
+        let payload_length = self.rlp_payload_length();
+        payload_length + length_of_length(payload_length)
+    }
+}
+
+impl<'a> LeafNodeRef<'a> {
+    /// Creates a new leaf node with the given key and value.
+    pub const fn new(key: &'a Nibbles, value: &'a [u8]) -> Self {
+        Self { key, value }
+    }
+
+    /// Returns the length of RLP encoded fields of leaf node.
+    fn rlp_payload_length(&self) -> usize {
+        let mut encoded_key_len = self.key.len() / 2 + 1;
+        // For leaf nodes the first byte cannot be greater than 0x80.
+        if encoded_key_len != 1 {
+            encoded_key_len += length_of_length(encoded_key_len);
+        }
+        encoded_key_len + Encodable::length(&self.value)
+    }
+
+    /// RLP encodes the node and returns either RLP(Node) or RLP(keccak(RLP(node)))
+    /// depending on if the serialized node was longer than a keccak).
+    pub fn rlp(&self, out: &mut Vec<u8>) -> Vec<u8> {
+        self.encode(out);
+        rlp_node(out)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -73,8 +146,9 @@ mod tests {
     fn rlp_leaf_node_roundtrip() {
         let nibble = Nibbles::from_nibbles_unchecked(hex!("0604060f"));
         let val = hex!("76657262");
-        let leaf = LeafNode::new(&nibble, &val);
-        let rlp = leaf.rlp(&mut vec![]);
+        let leaf = LeafNode::new(nibble, val.to_vec());
+        let rlp = leaf.as_ref().rlp(&mut vec![]);
         assert_eq!(rlp, hex!("c98320646f8476657262"));
+        assert_eq!(LeafNode::decode(&mut &rlp[..]).unwrap(), leaf);
     }
 }

src/nodes/mod.rs

@@ -14,7 +14,7 @@ mod extension;
 pub use extension::ExtensionNode;
 
 mod leaf;
-pub use leaf::LeafNode;
+pub use leaf::{LeafNode, LeafNodeRef};
 
 /// The range of valid child indexes.
 pub const CHILD_INDEX_RANGE: Range<u8> = 0..16;