feat(mpt): Refactor TrieNode (#172)

Refactors the `TrieNode` type to not contain the intermediate
`NodeElement`, but rather be recursive upon itself. This allows us to
create openings on a root `TrieNode` within the cache database and cache
all opened nodes throughout the execution of each block.

crates/mpt/Cargo.toml

@@ -13,10 +13,14 @@ homepage.workspace = true
 anyhow.workspace = true
 tracing.workspace = true
 alloy-primitives = { workspace = true, features = ["rlp"] }
-alloy-rlp = { workspace = true, default-features = false }
+alloy-rlp.workspace = true
+alloy-consensus.workspace = true
 
 # External
 alloy-trie = { version = "0.3.1", default-features = false }
+smallvec = "1.13"
+revm-primitives = { version = "3.1.1", default-features = false }
+revm = { version = "8.0.0", default-features = false }
 
 [dev-dependencies]
 alloy-consensus.workspace = true
@@ -26,3 +30,4 @@ alloy-rpc-types = { git = "https://github.com/alloy-rs/alloy", rev = "e3f2f07" }
 alloy-transport-http = { git = "https://github.com/alloy-rs/alloy", rev = "e3f2f07" }
 reqwest = "0.12"
 tracing-subscriber = "0.3.18"
+futures = { version = "0.3.30", default-features = false }

crates/mpt/README.md

@@ -1,3 +1,3 @@
 # `kona-mpt`
 
-Utilities for interacting with and iterating through a merkle patricia trie
+Utilities for interacting with a merkle patricia trie in the client program.

crates/mpt/src/lib.rs

@@ -7,7 +7,7 @@
 extern crate alloc;
 
 mod node;
-pub use node::{NodeElement, TrieNode};
+pub use node::TrieNode;
 
 mod list_walker;
 pub use list_walker::OrderedListWalker;

crates/mpt/src/list_walker.rs

@@ -1,12 +1,12 @@
 //! This module contains the [OrderedListWalker] struct, which allows for traversing an MPT root of
 //! a derivable ordered list.
 
-use crate::{NodeElement, TrieNode};
+use crate::TrieNode;
 use alloc::{collections::VecDeque, vec};
 use alloy_primitives::{Bytes, B256};
 use alloy_rlp::{Decodable, EMPTY_STRING_CODE};
 use anyhow::{anyhow, Result};
-use core::{fmt::Display, marker::PhantomData};
+use core::marker::PhantomData;
 
 /// A [OrderedListWalker] allows for traversing over a Merkle Patricia Trie containing a derivable
 /// ordered list.
@@ -19,7 +19,7 @@ pub struct OrderedListWalker<PreimageFetcher> {
     root: B256,
     /// The leaf nodes of the derived list, in order. [None] if the tree has yet to be fully
     /// traversed with [Self::hydrate].
-    inner: Option<VecDeque<Bytes>>,
+    inner: Option<VecDeque<(Bytes, Bytes)>>,
     /// Phantom data
     _phantom: PhantomData<PreimageFetcher>,
 }
@@ -55,7 +55,7 @@ where
         // With small lists the iterator seems to use 0x80 (RLP empty string, unlike the others)
         // as key for item 0, causing it to come last. We need to account for this, pulling the
         // first element into its proper position.
-        let mut ordered_list = Self::fetch_leaves(root_trie_node, fetcher)?;
+        let mut ordered_list = Self::fetch_leaves(&root_trie_node, fetcher)?;
         if !ordered_list.is_empty() {
             if ordered_list.len() <= EMPTY_STRING_CODE as usize {
                 // If the list length is < 0x80, the final element is the first element.
@@ -74,53 +74,67 @@ where
         Ok(())
     }
 
+    /// Takes the inner list of the [OrderedListWalker], returning it and setting the inner list to
+    /// [None].
+    pub fn take_inner(&mut self) -> Option<VecDeque<(Bytes, Bytes)>> {
+        self.inner.take()
+    }
+
     /// Traverses a [TrieNode], returning all values of child [TrieNode::Leaf] variants.
-    fn fetch_leaves(trie_node: TrieNode, fetcher: PreimageFetcher) -> Result<VecDeque<Bytes>> {
+    fn fetch_leaves(
+        trie_node: &TrieNode,
+        fetcher: PreimageFetcher,
+    ) -> Result<VecDeque<(Bytes, Bytes)>> {
         match trie_node {
             TrieNode::Branch { stack } => {
                 let mut leaf_values = VecDeque::with_capacity(stack.len());
-                for item in stack.into_iter() {
+                for item in stack.iter() {
                     match item {
-                        NodeElement::String(s) => {
+                        TrieNode::Blinded { commitment } => {
                             // If the string is a hash, we need to grab the preimage for it and
                             // continue recursing.
-                            let trie_node = Self::get_trie_node(s.as_ref(), fetcher)?;
-                            leaf_values.append(&mut Self::fetch_leaves(trie_node, fetcher)?);
+                            let trie_node = Self::get_trie_node(commitment.as_ref(), fetcher)?;
+                            leaf_values.append(&mut Self::fetch_leaves(&trie_node, fetcher)?);
+                        }
+                        TrieNode::Empty => { /* Skip over empty nodes, we're looking for values. */
                         }
-                        list @ NodeElement::List(_) => {
-                            let trie_node = list.try_list_into_node()?;
-                            leaf_values.append(&mut Self::fetch_leaves(trie_node, fetcher)?);
+                        item => {
+                            // If the item is already retrieved, recurse on it.
+                            leaf_values.append(&mut Self::fetch_leaves(item, fetcher)?);
                         }
-                        _ => { /* Skip over empty lists and strings; We're looking for leaves */ }
                     }
                 }
                 Ok(leaf_values)
             }
-            TrieNode::Leaf { value, .. } => Ok(vec![value].into()),
+            TrieNode::Leaf { key, value } => Ok(vec![(key.clone(), value.clone())].into()),
             TrieNode::Extension { node, .. } => {
                 // If the node is a hash, we need to grab the preimage for it and continue
-                // recursing.
-                let trie_node = Self::get_trie_node(node.as_ref(), fetcher)?;
-                Ok(Self::fetch_leaves(trie_node, fetcher)?)
+                // recursing. If it is already retrieved, recurse on it.
+                match node.as_ref() {
+                    TrieNode::Blinded { commitment } => {
+                        let trie_node = Self::get_trie_node(commitment.as_ref(), fetcher)?;
+                        Ok(Self::fetch_leaves(&trie_node, fetcher)?)
+                    }
+                    node => Ok(Self::fetch_leaves(node, fetcher)?),
+                }
             }
+            _ => anyhow::bail!("Invalid trie node type encountered"),
         }
     }
 
     /// Grabs the preimage of `hash` using `fetcher`, and attempts to decode the preimage data into
     /// a [TrieNode]. Will error if the conversion of `T` into [B256] fails.
     fn get_trie_node<T>(hash: T, fetcher: PreimageFetcher) -> Result<TrieNode>
     where
-        T: TryInto<B256>,
-        <T as TryInto<B256>>::Error: Display,
+        T: Into<B256>,
     {
-        let hash = hash.try_into().map_err(|e| anyhow!("Error in conversion: {e}"))?;
-        let preimage = fetcher(hash)?;
+        let preimage = fetcher(hash.into())?;
         TrieNode::decode(&mut preimage.as_ref()).map_err(|e| anyhow!(e))
     }
 }
 
 impl<PreimageFetcher> Iterator for OrderedListWalker<PreimageFetcher> {
-    type Item = Bytes;
+    type Item = (Bytes, Bytes);
 
     fn next(&mut self) -> Option<Self::Item> {
         match self.inner {
@@ -152,7 +166,7 @@ mod test {
 
         assert_eq!(
             list.into_iter()
-                .map(|rlp| ReceiptEnvelope::decode_2718(&mut rlp.as_ref()).unwrap())
+                .map(|(_, rlp)| ReceiptEnvelope::decode_2718(&mut rlp.as_ref()).unwrap())
                 .collect::<Vec<_>>(),
             envelopes
         );
@@ -167,7 +181,7 @@ mod test {
 
         assert_eq!(
             list.into_iter()
-                .map(|rlp| TxEnvelope::decode(&mut rlp.as_ref()).unwrap())
+                .map(|(_, rlp)| TxEnvelope::decode(&mut rlp.as_ref()).unwrap())
                 .collect::<Vec<_>>(),
             envelopes
         );
@@ -194,7 +208,7 @@ mod test {
             list.inner
                 .unwrap()
                 .iter()
-                .map(|v| String::decode(&mut v.as_ref()).unwrap())
+                .map(|(_, v)| { String::decode(&mut v.as_ref()).unwrap() })
                 .collect::<Vec<_>>(),
             VALUES
         );