feat(mpt): TrieNode retrieval

Adds a function to the trie node that allows for opening the trie at a
given path. As the trie is opened, intermediate nodes are persisted
within the trie node passed.

crates/mpt/src/node.rs

@@ -4,6 +4,7 @@
 use alloc::{boxed::Box, collections::VecDeque, vec, vec::Vec};
 use alloy_primitives::{keccak256, Bytes, B256};
 use alloy_rlp::{Buf, BufMut, Decodable, Encodable, Header, EMPTY_STRING_CODE};
+use alloy_trie::Nibbles;
 use anyhow::{anyhow, Result};
 
 /// The length of the branch list when RLP encoded
@@ -115,6 +116,98 @@ impl TrieNode {
             self
         }
     }
+
+    /// Walks down the trie to a leaf value with the given key, if it exists. Preimages for blinded
+    /// nodes along the path are fetched using the `fetcher` function, and persisted in the inner
+    /// [TrieNode] elements.
+    ///
+    /// ## Takes
+    /// - `item_key` - The nibbles representation of the key being retrieved
+    /// - `trie_node` - The root trie node
+    /// - `pos` - The number of nibbles that have already been traversed in the `item_key`
+    /// - `fetcher` - The preimage fetcher for intermediate blinded nodes
+    ///
+    /// ## Returns
+    /// - `Err(_)` - Could not retrieve the node with the given key from the trie.
+    /// - `Ok((_, _))` - The key and value of the node
+    pub fn open<'a>(
+        &'a mut self,
+        item_key: &Nibbles,
+        mut pos: usize,
+        fetcher: impl Fn(B256) -> Result<Bytes> + Copy,
+    ) -> Result<&'a Bytes> {
+        match self {
+            TrieNode::Branch { ref mut stack } => {
+                let branch_nibble = item_key[pos] as usize;
+                pos += 1;
+
+                let el =
+                    stack.get_mut(branch_nibble).ok_or(anyhow!("Key does not exist in trie"))?;
+                match el {
+                    TrieNode::Blinded { commitment } => {
+                        // If the string is a hash, we need to grab the preimage for it and
+                        // continue recursing.
+                        let trie_node = TrieNode::decode(&mut fetcher(*commitment)?.as_ref())
+                            .map_err(|e| anyhow!(e))?;
+                        *el = trie_node;
+
+                        // If the value was found in the blinded node, return it.
+                        if let Ok(value) = el.open(item_key, pos, fetcher) {
+                            return Ok(value);
+                        }
+                    }
+                    TrieNode::Empty => anyhow::bail!("Key does not exist in trie"),
+                    node => {
+                        // If the value was found in the blinded node, return it.
+                        if let Ok(value) = node.open(item_key, pos, fetcher) {
+                            return Ok(value);
+                        }
+                    }
+                };
+
+                anyhow::bail!("Key does not exist in trie");
+            }
+            TrieNode::Leaf { key, value } => {
+                let key_nibbles = Nibbles::unpack(key.clone());
+                let shared_nibbles = key_nibbles[1..].as_ref();
+
+                // If the key length is one, it only contains the prefix and no shared nibbles.
+                // Return the key and value.
+                if key.len() == 1 || pos + shared_nibbles.len() >= item_key.len() {
+                    return Ok(value);
+                }
+
+                let item_key_nibbles = item_key[pos..pos + shared_nibbles.len()].as_ref();
+
+                if item_key_nibbles == shared_nibbles {
+                    Ok(value)
+                } else {
+                    anyhow::bail!("Key does not exist in trie");
+                }
+            }
+            TrieNode::Extension { prefix, node } => {
+                let prefix_nibbles = Nibbles::unpack(prefix);
+                let shared_nibbles = prefix_nibbles[1..].as_ref();
+                let item_key_nibbles = item_key[pos..pos + shared_nibbles.len()].as_ref();
+                if item_key_nibbles == shared_nibbles {
+                    // Increase the offset within the key by the length of the shared nibbles
+                    pos += shared_nibbles.len();
+
+                    // Follow extension branch
+                    if let TrieNode::Blinded { commitment } = node.as_ref() {
+                        *node = Box::new(
+                            TrieNode::decode(&mut fetcher(*commitment)?.as_ref())
+                                .map_err(|e| anyhow!(e))?,
+                        );
+                    }
+                    node.open(item_key, pos, fetcher)
+                } else {
+                    anyhow::bail!("Key does not exist in trie");
+                }
+            }
+            _ => anyhow::bail!("Invalid trie node type encountered"),
+        }
+    }
 }
 
 impl Encodable for TrieNode {
@@ -261,8 +354,12 @@ fn rlp_list_element_length(buf: &mut &[u8]) -> alloy_rlp::Result<usize> {
 #[cfg(test)]
 mod test {
     use super::*;
-    use alloc::vec;
-    use alloy_primitives::{b256, bytes, hex};
+    use crate::{test_util::ordered_trie_with_encoder, TrieNode};
+    use alloc::{collections::BTreeMap, vec, vec::Vec};
+    use alloy_primitives::{b256, bytes, hex, keccak256, Bytes, B256};
+    use alloy_rlp::{Decodable, Encodable, EMPTY_STRING_CODE};
+    use alloy_trie::Nibbles;
+    use anyhow::{anyhow, Result};
 
     #[test]
     fn test_decode_branch() {
@@ -322,7 +419,7 @@ mod test {
             hex!("e58300646fa0f3fe8b3c5b21d3e52860f1e4a5825a6100bb341069c1e88f4ebf6bd98de0c190");
         let mut rlp_buf = Vec::new();
 
-        let opened = TrieNode::Leaf { key: bytes!("30"), value: bytes!("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF") };
+        let opened = TrieNode::Leaf { key: bytes!("30"), value: [0xFF; 64].into() };
         opened.encode(&mut rlp_buf);
         let blinded = TrieNode::Blinded { commitment: keccak256(&rlp_buf) };
 
@@ -341,4 +438,37 @@ mod test {
         let expected = TrieNode::Leaf { key: bytes!("20646f"), value: bytes!("76657262FF") };
         assert_eq!(expected, TrieNode::decode(&mut LEAF_RLP.as_slice()).unwrap());
     }
+
+    #[test]
+    fn test_retrieve_from_trie_simple() {
+        const VALUES: [&str; 5] = ["yeah", "dog", ", ", "laminar", "flow"];
+
+        let mut trie = ordered_trie_with_encoder(&VALUES, |v, buf| v.encode(buf));
+        let root = trie.root();
+
+        let preimages =
+            trie.take_proofs().into_iter().fold(BTreeMap::default(), |mut acc, (_, value)| {
+                acc.insert(keccak256(value.as_ref()), value);
+                acc
+            });
+        let fetcher = |h: B256| -> Result<Bytes> {
+            preimages.get(&h).cloned().ok_or(anyhow!("Failed to find preimage"))
+        };
+
+        let mut root_node = TrieNode::decode(&mut fetcher(root).unwrap().as_ref()).unwrap();
+        for (i, value) in VALUES.iter().enumerate() {
+            let key_nibbles = Nibbles::unpack([if i == 0 { EMPTY_STRING_CODE } else { i as u8 }]);
+            let v = root_node.open(&key_nibbles, 0, fetcher).unwrap();
+
+            let mut encoded_value = Vec::with_capacity(value.length());
+            value.encode(&mut encoded_value);
+
+            assert_eq!(v, encoded_value.as_slice());
+        }
+
+        let TrieNode::Blinded { commitment } = root_node.blind() else {
+            panic!("Expected blinded root node");
+        };
+        assert_eq!(commitment, root);
+    }
 }