fix: use vector of arbitrary length for leaf node value (#56)

* fix: use vector of arbitrary length for leaf node value

* comments

* to_vec -> into

src/hash_builder/mod.rs

@@ -579,8 +579,8 @@ mod tests {
     #[test]
     fn manual_branch_node_ok() {
         let raw_input = vec![
-            (hex!("646f").to_vec(), RlpNode::from_raw(&hex!("76657262")).unwrap()),
-            (hex!("676f6f64").to_vec(), RlpNode::from_raw(&hex!("7075707079")).unwrap()),
+            (hex!("646f").to_vec(), hex!("76657262").to_vec()),
+            (hex!("676f6f64").to_vec(), hex!("7075707079").to_vec()),
         ];
         let expected = triehash_trie_root(raw_input.clone());
 

src/nodes/branch.rs

@@ -334,10 +334,7 @@ mod tests {
         let encoded = alloy_rlp::encode(&sparse_node);
         assert_eq!(BranchNode::decode(&mut &encoded[..]).unwrap(), sparse_node);
 
-        let leaf_child = LeafNode::new(
-            Nibbles::from_nibbles(hex!("0203")),
-            RlpNode::from_raw(&hex!("1234")).unwrap(),
-        );
+        let leaf_child = LeafNode::new(Nibbles::from_nibbles(hex!("0203")), hex!("1234").to_vec());
         let mut buf = vec![];
         let leaf_rlp = leaf_child.as_ref().rlp(&mut buf);
         let branch_with_leaf = BranchNode::new(vec![leaf_rlp.clone()], TrieMask::new(0b0010));

src/nodes/leaf.rs

@@ -18,7 +18,7 @@ pub struct LeafNode {
     /// The key for this leaf node.
     pub key: Nibbles,
     /// The node value.
-    pub value: RlpNode,
+    pub value: Vec<u8>,
 }
 
 impl fmt::Debug for LeafNode {
@@ -58,7 +58,7 @@ impl Decodable for LeafNode {
         };
 
         let key = unpack_path_to_nibbles(first, &encoded_key[1..]);
-        let value = RlpNode::decode(&mut bytes)?;
+        let value = Bytes::decode(&mut bytes)?.into();
         Ok(Self { key, value })
     }
 }
@@ -71,7 +71,7 @@ impl LeafNode {
     pub const ODD_FLAG: u8 = 0x30;
 
     /// Creates a new leaf node with the given key and value.
-    pub const fn new(key: Nibbles, value: RlpNode) -> Self {
+    pub const fn new(key: Nibbles, value: Vec<u8>) -> Self {
         Self { key, value }
     }
 
@@ -155,7 +155,7 @@ mod tests {
     fn rlp_leaf_node_roundtrip() {
         let nibble = Nibbles::from_nibbles_unchecked(hex!("0604060f"));
         let val = hex!("76657262");
-        let leaf = LeafNode::new(nibble, RlpNode::from_raw(&val).unwrap());
+        let leaf = LeafNode::new(nibble, val.to_vec());
         let rlp = leaf.as_ref().rlp(&mut vec![]);
         assert_eq!(rlp.as_ref(), hex!("c98320646f8476657262"));
         assert_eq!(LeafNode::decode(&mut &rlp[..]).unwrap(), leaf);

src/nodes/mod.rs

@@ -1,6 +1,6 @@
 //! Various branch nodes produced by the hash builder.
 
-use alloy_primitives::B256;
+use alloy_primitives::{Bytes, B256};
 use alloy_rlp::{Decodable, Encodable, Header, EMPTY_STRING_CODE};
 use core::ops::Range;
 use nybbles::Nibbles;
@@ -112,7 +112,8 @@ impl Decodable for TrieNode {
 
                 let key = unpack_path_to_nibbles(first, &encoded_key[1..]);
                 let node = if key_flag == LeafNode::EVEN_FLAG || key_flag == LeafNode::ODD_FLAG {
-                    Self::Leaf(LeafNode::new(key, RlpNode::decode(&mut items.remove(0))?))
+                    let value = Bytes::decode(&mut items.remove(0))?.into();
+                    Self::Leaf(LeafNode::new(key, value))
                 } else {
                     // We don't decode value because it is expected to be RLP encoded.
                     Self::Extension(ExtensionNode::new(
@@ -270,7 +271,7 @@ mod tests {
     fn rlp_zero_value_leaf_roundtrip() {
         let leaf = TrieNode::Leaf(LeafNode::new(
             Nibbles::from_nibbles_unchecked(hex!("0604060f")),
-            RlpNode::from_raw(&alloy_rlp::encode(alloy_primitives::U256::ZERO)).unwrap(),
+            alloy_rlp::encode(alloy_primitives::U256::ZERO),
         ));
         let rlp = leaf.rlp(&mut vec![]);
         assert_eq!(rlp[..], hex!("c68320646f8180"));
@@ -282,7 +283,7 @@ mod tests {
         // leaf
         let leaf = TrieNode::Leaf(LeafNode::new(
             Nibbles::from_nibbles_unchecked(hex!("0604060f")),
-            RlpNode::from_raw(&hex!("76657262")).unwrap(),
+            hex!("76657262").to_vec(),
         ));
         let rlp = leaf.rlp(&mut vec![]);
         assert_eq!(rlp[..], hex!("c98320646f8476657262"));

src/proof/verify.rs

@@ -1,5 +1,7 @@
 //! Proof verification logic.
 
+use core::ops::Deref;
+
 use crate::{
     nodes::{BranchNode, RlpNode, TrieNode, CHILD_INDEX_RANGE},
     proof::ProofVerificationError,
@@ -17,71 +19,100 @@ use nybbles::Nibbles;
 pub fn verify_proof<'a, I>(
     root: B256,
     key: Nibbles,
-    value: Option<Vec<u8>>,
+    expected_value: Option<Vec<u8>>,
     proof: I,
 ) -> Result<(), ProofVerificationError>
 where
     I: IntoIterator<Item = &'a Bytes>,
 {
     let mut proof = proof.into_iter().peekable();
 
+    // If the proof is empty or contains only an empty node, the expected value must be None.
     if proof.peek().map_or(true, |node| node.as_ref() == [EMPTY_STRING_CODE]) {
         return if root == EMPTY_ROOT_HASH {
-            if value.is_none() {
+            if expected_value.is_none() {
                 Ok(())
             } else {
                 Err(ProofVerificationError::ValueMismatch {
                     path: key,
                     got: None,
-                    expected: value.map(Bytes::from),
+                    expected: expected_value.map(Bytes::from),
                 })
             }
         } else {
             Err(ProofVerificationError::RootMismatch { got: EMPTY_ROOT_HASH, expected: root })
         };
     }
 
-    let mut walked_path = Nibbles::default();
-    let mut next_value = Some(RlpNode::word_rlp(&root));
+    let mut walked_path = Nibbles::with_capacity(key.len());
+    let mut last_decoded_node = Some(NodeDecodingResult::Node(RlpNode::word_rlp(&root)));
     for node in proof {
-        if Some(RlpNode::from_rlp(node)) != next_value {
+        // Check if the node that we just decoded (or root node, if we just started) matches
+        // the expected node from the proof.
+        if Some(RlpNode::from_rlp(node).as_slice()) != last_decoded_node.as_deref() {
             let got = Some(Bytes::copy_from_slice(node));
-            let expected = next_value.map(|b| Bytes::copy_from_slice(&b));
+            let expected = last_decoded_node.as_deref().map(Bytes::copy_from_slice);
             return Err(ProofVerificationError::ValueMismatch { path: walked_path, got, expected });
         }
 
-        next_value = match TrieNode::decode(&mut &node[..])? {
+        // Decode the next node from the proof.
+        last_decoded_node = match TrieNode::decode(&mut &node[..])? {
             TrieNode::Branch(branch) => process_branch(branch, &mut walked_path, &key)?,
             TrieNode::Extension(extension) => {
                 walked_path.extend_from_slice(&extension.key);
-                Some(extension.child)
+                Some(NodeDecodingResult::Node(extension.child))
             }
             TrieNode::Leaf(leaf) => {
                 walked_path.extend_from_slice(&leaf.key);
-                Some(leaf.value)
+                Some(NodeDecodingResult::Value(leaf.value))
             }
             TrieNode::EmptyRoot => return Err(ProofVerificationError::UnexpectedEmptyRoot),
         };
     }
 
-    next_value = next_value.filter(|_| walked_path == key);
-    if next_value.as_deref() == value.as_deref() {
+    // Last decoded node should have the key that we are looking for.
+    last_decoded_node = last_decoded_node.filter(|_| walked_path == key);
+    if last_decoded_node.as_deref() == expected_value.as_deref() {
         Ok(())
     } else {
         Err(ProofVerificationError::ValueMismatch {
             path: key,
-            got: next_value.as_deref().map(Vec::from).map(Bytes::from),
-            expected: value.map(Bytes::from),
+            got: last_decoded_node.as_deref().map(Bytes::copy_from_slice),
+            expected: expected_value.map(Bytes::from),
         })
     }
 }
 
+/// The result of decoding a node from the proof.
+///
+/// - [`TrieNode::Branch`] is decoded into a [`NodeDecodingResult::Value`] if the node at the
+///   specified nibble was decoded into an in-place encoded [`TrieNode::Leaf`], or into a
+///   [`NodeDecodingResult::Node`] otherwise.
+/// - [`TrieNode::Extension`] is always decoded into a [`NodeDecodingResult::Node`].
+/// - [`TrieNode::Leaf`] is always decoded into a [`NodeDecodingResult::Value`].
+#[derive(Debug, PartialEq, Eq)]
+enum NodeDecodingResult {
+    Node(RlpNode),
+    Value(Vec<u8>),
+}
+
+impl Deref for NodeDecodingResult {
+    type Target = [u8];
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            NodeDecodingResult::Node(node) => node.as_slice(),
+            NodeDecodingResult::Value(value) => value,
+        }
+    }
+}
+
 #[inline]
 fn process_branch(
     mut branch: BranchNode,
     walked_path: &mut Nibbles,
     key: &Nibbles,
-) -> Result<Option<RlpNode>, ProofVerificationError> {
+) -> Result<Option<NodeDecodingResult>, ProofVerificationError> {
     if let Some(next) = key.get(walked_path.len()) {
         let mut stack_ptr = branch.as_ref().first_child_index();
         for index in CHILD_INDEX_RANGE {
@@ -91,7 +122,7 @@ fn process_branch(
 
                     let child = branch.stack.remove(stack_ptr);
                     if child.len() == B256::len_bytes() + 1 {
-                        return Ok(Some(child));
+                        return Ok(Some(NodeDecodingResult::Node(child)));
                     } else {
                         // This node is encoded in-place.
                         match TrieNode::decode(&mut &child[..])? {
@@ -106,12 +137,12 @@ fn process_branch(
                                 walked_path.extend_from_slice(&child_extension.key);
 
                                 // If the extension node's child is a hash, the encoded extension
-                                // node itself wouldn't fit for encoding in- place. So this
-                                // extension node must have a child that is also encoded in-place.
+                                // node itself wouldn't fit for encoding in-place. So this extension
+                                // node must have a child that is also encoded in-place.
                                 //
                                 // Since the child cannot be a leaf node (otherwise this node itself
-                                // is a leaf node to begin with, the child must also be a branch
-                                // encoded in-place.
+                                // would be a leaf node, not an extension node), the child must be a
+                                // branch node encoded in-place.
                                 match TrieNode::decode(&mut &child_extension.child[..])? {
                                     TrieNode::Branch(extension_child_branch) => {
                                         return process_branch(
@@ -129,7 +160,7 @@ fn process_branch(
                             }
                             TrieNode::Leaf(child_leaf) => {
                                 walked_path.extend_from_slice(&child_leaf.key);
-                                return Ok(Some(child_leaf.value));
+                                return Ok(Some(NodeDecodingResult::Value(child_leaf.value)));
                             }
                             TrieNode::EmptyRoot => {
                                 return Err(ProofVerificationError::UnexpectedEmptyRoot)
@@ -467,7 +498,7 @@ mod tests {
 
         let mut buffer = vec![];
 
-        let value = RlpNode::from_raw(&[0x64]).unwrap();
+        let value = vec![0x64];
         let child_leaf = TrieNode::Leaf(LeafNode::new(Nibbles::from_nibbles([0xa]), value.clone()));
 
         let child_branch = TrieNode::Branch(BranchNode::new(