feat(TrieProvider): Abstract TrieNode retrieval (#787)

op-rs · Nov 6, 2024 · 6106301 · 6106301
1 parent 7c5567b
commit 6106301
Show file tree

Hide file tree

Showing 11 changed files with 89 additions and 46 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/bin/client/src/l1/chain_provider.rs b/bin/client/src/l1/chain_provider.rs
@@ -8,7 +8,7 @@ use alloy_primitives::{Bytes, B256};
 use alloy_rlp::Decodable;
 use async_trait::async_trait;
 use kona_derive::traits::ChainProvider;
-use kona_mpt::{OrderedListWalker, TrieProvider};
+use kona_mpt::{OrderedListWalker, TrieNode, TrieProvider};
 use kona_preimage::{CommsClient, PreimageKey, PreimageKeyType};
 use op_alloy_protocol::BlockInfo;
 
@@ -138,15 +138,19 @@ impl<T: CommsClient + Sync + Send> ChainProvider for OracleL1ChainProvider<T> {
 impl<T: CommsClient> TrieProvider for OracleL1ChainProvider<T> {
     type Error = OracleProviderError;
 
-    fn trie_node_preimage(&self, key: B256) -> Result<Bytes, Self::Error> {
+    fn trie_node_by_hash(&self, key: B256) -> Result<TrieNode, Self::Error> {
         // On L1, trie node preimages are stored as keccak preimage types in the oracle. We assume
         // that a hint for these preimages has already been sent, prior to this call.
         kona_common::block_on(async move {
-            self.oracle
-                .get(PreimageKey::new(*key, PreimageKeyType::Keccak256))
-                .await
-                .map(Into::into)
-                .map_err(OracleProviderError::Preimage)
+            TrieNode::decode(
+                &mut self
+                    .oracle
+                    .get(PreimageKey::new(*key, PreimageKeyType::Keccak256))
+                    .await
+                    .map_err(OracleProviderError::Preimage)?
+                    .as_ref(),
+            )
+            .map_err(OracleProviderError::Rlp)
         })
     }
 

diff --git a/bin/client/src/l2/chain_provider.rs b/bin/client/src/l2/chain_provider.rs
@@ -8,7 +8,7 @@ use alloy_primitives::{Address, Bytes, B256};
 use alloy_rlp::Decodable;
 use async_trait::async_trait;
 use kona_derive::traits::L2ChainProvider;
-use kona_mpt::{OrderedListWalker, TrieHinter, TrieProvider};
+use kona_mpt::{OrderedListWalker, TrieHinter, TrieNode, TrieProvider};
 use kona_preimage::{CommsClient, PreimageKey, PreimageKeyType};
 use op_alloy_consensus::{OpBlock, OpTxEnvelope};
 use op_alloy_genesis::{RollupConfig, SystemConfig};
@@ -144,15 +144,19 @@ impl<T: CommsClient + Send + Sync> L2ChainProvider for OracleL2ChainProvider<T>
 impl<T: CommsClient> TrieProvider for OracleL2ChainProvider<T> {
     type Error = OracleProviderError;
 
-    fn trie_node_preimage(&self, key: B256) -> Result<Bytes, OracleProviderError> {
+    fn trie_node_by_hash(&self, key: B256) -> Result<TrieNode, OracleProviderError> {
         // On L2, trie node preimages are stored as keccak preimage types in the oracle. We assume
         // that a hint for these preimages has already been sent, prior to this call.
         kona_common::block_on(async move {
-            self.oracle
-                .get(PreimageKey::new(*key, PreimageKeyType::Keccak256))
-                .await
-                .map(Into::into)
-                .map_err(OracleProviderError::Preimage)
+            TrieNode::decode(
+                &mut self
+                    .oracle
+                    .get(PreimageKey::new(*key, PreimageKeyType::Keccak256))
+                    .await
+                    .map_err(OracleProviderError::Preimage)?
+                    .as_ref(),
+            )
+            .map_err(OracleProviderError::Rlp)
         })
     }
 

diff --git a/crates/executor/benches/execution.rs b/crates/executor/benches/execution.rs
@@ -8,7 +8,7 @@ use alloy_rpc_types_engine::PayloadAttributes;
 use anyhow::{anyhow, Result};
 use criterion::{criterion_group, criterion_main, Bencher, Criterion};
 use kona_executor::StatelessL2BlockExecutor;
-use kona_mpt::{NoopTrieHinter, TrieProvider};
+use kona_mpt::{NoopTrieHinter, TrieNode, TrieProvider};
 use op_alloy_genesis::{RollupConfig, OP_MAINNET_BASE_FEE_PARAMS};
 use op_alloy_rpc_types_engine::OpPayloadAttributes;
 use pprof::criterion::{Output, PProfProfiler};
@@ -37,11 +37,16 @@ impl TestdataTrieProvider {
 impl TrieProvider for TestdataTrieProvider {
     type Error = anyhow::Error;
 
-    fn trie_node_preimage(&self, key: B256) -> Result<Bytes> {
-        self.preimages
-            .get(&key)
-            .cloned()
-            .ok_or_else(|| anyhow!("Preimage not found for key: {}", key))
+    fn trie_node_by_hash(&self, key: B256) -> Result<TrieNode> {
+        TrieNode::decode(
+            &mut self
+                .preimages
+                .get(&key)
+                .cloned()
+                .ok_or_else(|| anyhow!("Preimage not found for key: {}", key))?
+                .as_ref(),
+        )
+        .map_err(Into::into)
     }
 
     fn bytecode_by_hash(&self, code_hash: B256) -> Result<Bytes> {

diff --git a/crates/executor/src/executor/mod.rs b/crates/executor/src/executor/mod.rs
@@ -456,7 +456,7 @@ mod test {
     use alloy_rlp::Decodable;
     use alloy_rpc_types_engine::PayloadAttributes;
     use anyhow::{anyhow, Result};
-    use kona_mpt::NoopTrieHinter;
+    use kona_mpt::{NoopTrieHinter, TrieNode};
     use op_alloy_genesis::OP_MAINNET_BASE_FEE_PARAMS;
     use serde::Deserialize;
     use std::collections::HashMap;
@@ -483,11 +483,16 @@ mod test {
     impl TrieProvider for TestdataTrieProvider {
         type Error = anyhow::Error;
 
-        fn trie_node_preimage(&self, key: B256) -> Result<Bytes> {
-            self.preimages
-                .get(&key)
-                .cloned()
-                .ok_or_else(|| anyhow!("Preimage not found for key: {}", key))
+        fn trie_node_by_hash(&self, key: B256) -> Result<TrieNode> {
+            TrieNode::decode(
+                &mut self
+                    .preimages
+                    .get(&key)
+                    .cloned()
+                    .ok_or_else(|| anyhow!("Preimage not found for key: {}", key))?
+                    .as_ref(),
+            )
+            .map_err(Into::into)
         }
 
         fn bytecode_by_hash(&self, code_hash: B256) -> Result<Bytes> {

diff --git a/crates/mpt/Cargo.toml b/crates/mpt/Cargo.toml
@@ -14,6 +14,7 @@ workspace = true
 [dependencies]
 # General
 derive_more = { workspace = true, features = ["full"] }
+serde = { workspace = true, optional = true, features = ["derive", "alloc"] }
 
 # Revm + Alloy
 alloy-rlp.workspace = true
@@ -38,6 +39,14 @@ criterion = { workspace = true, features = ["html_reports"] }
 tracing-subscriber = { workspace = true, features = ["fmt"] }
 pprof = { workspace = true, features = ["criterion", "flamegraph", "frame-pointer"] }
 
+[features]
+default = ["serde"]
+serde = [
+    "dep:serde",
+    "alloy-primitives/serde",
+    "alloy-trie/serde"
+]
+
 [[bench]]
 name = "trie_node"
 harness = false
diff --git a/crates/mpt/src/list_walker.rs b/crates/mpt/src/list_walker.rs
@@ -7,7 +7,7 @@ use crate::{
 };
 use alloc::{collections::VecDeque, string::ToString, vec};
 use alloy_primitives::{Bytes, B256};
-use alloy_rlp::{Decodable, EMPTY_STRING_CODE};
+use alloy_rlp::EMPTY_STRING_CODE;
 use core::marker::PhantomData;
 
 /// A [OrderedListWalker] allows for traversing over a Merkle Patricia Trie containing a derivable
@@ -132,11 +132,9 @@ where
     where
         T: Into<B256>,
     {
-        let preimage = fetcher
-            .trie_node_preimage(hash.into())
-            .map_err(|e| TrieNodeError::Provider(e.to_string()))?;
-        TrieNode::decode(&mut preimage.as_ref())
-            .map_err(TrieNodeError::RLPError)
+        fetcher
+            .trie_node_by_hash(hash.into())
+            .map_err(|e| TrieNodeError::Provider(e.to_string()))
             .map_err(Into::into)
     }
 }
@@ -176,7 +174,7 @@ mod test {
     use alloy_consensus::{ReceiptEnvelope, TxEnvelope};
     use alloy_primitives::keccak256;
     use alloy_provider::network::eip2718::Decodable2718;
-    use alloy_rlp::Encodable;
+    use alloy_rlp::{Decodable, Encodable};
 
     #[tokio::test]
     async fn test_online_list_walker_receipts() {

diff --git a/crates/mpt/src/node.rs b/crates/mpt/src/node.rs
@@ -62,6 +62,7 @@ const NIBBLE_WIDTH: usize = 4;
 /// to behave correctly if confronted with keys of varying lengths. Namely, this is because it does
 /// not support the `value` field in branch nodes, just like the Ethereum Merkle Patricia Trie.
 #[derive(Debug, Clone, Eq, PartialEq, Display)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub enum TrieNode {
     /// An empty [TrieNode] is represented as an [EMPTY_STRING_CODE] (0x80).
     Empty,
@@ -141,10 +142,9 @@ impl TrieNode {
                 // reach out to the fetcher.
                 *self = Self::Empty;
             } else {
-                let rlp = fetcher
-                    .trie_node_preimage(*commitment)
+                *self = fetcher
+                    .trie_node_by_hash(*commitment)
                     .map_err(|e| TrieNodeError::Provider(e.to_string()))?;
-                *self = Self::decode(&mut rlp.as_ref()).map_err(TrieNodeError::RLPError)?;
             }
         }
         Ok(())
@@ -794,8 +794,7 @@ mod test {
         );
         let fetcher = TrieNodeProvider::new(preimages, Default::default(), Default::default());
 
-        let mut root_node =
-            TrieNode::decode(&mut fetcher.trie_node_preimage(root).unwrap().as_ref()).unwrap();
+        let mut root_node = fetcher.trie_node_by_hash(root).unwrap();
         for (i, value) in VALUES.iter().enumerate() {
             let path_nibbles = Nibbles::unpack([if i == 0 { EMPTY_STRING_CODE } else { i as u8 }]);
             let v = root_node.open(&path_nibbles, &fetcher).unwrap().unwrap();

diff --git a/crates/mpt/src/noop.rs b/crates/mpt/src/noop.rs
@@ -1,7 +1,7 @@
 //! Trait implementations for `kona-mpt` traits that are effectively a no-op.
 //! Providers trait implementations for downstream users who do not require hinting.
 
-use crate::{TrieHinter, TrieProvider};
+use crate::{TrieHinter, TrieNode, TrieProvider};
 use alloc::string::String;
 use alloy_consensus::Header;
 use alloy_primitives::{Address, Bytes, B256, U256};
@@ -13,8 +13,8 @@ pub struct NoopTrieProvider;
 impl TrieProvider for NoopTrieProvider {
     type Error = String;
 
-    fn trie_node_preimage(&self, _key: B256) -> Result<Bytes, Self::Error> {
-        Ok(Bytes::new())
+    fn trie_node_by_hash(&self, _key: B256) -> Result<TrieNode, Self::Error> {
+        Ok(TrieNode::Empty)
     }
 
     fn bytecode_by_hash(&self, _code_hash: B256) -> Result<Bytes, Self::Error> {

diff --git a/crates/mpt/src/test_util.rs b/crates/mpt/src/test_util.rs
@@ -1,10 +1,11 @@
 //! Testing utilities for `kona-mpt`
 
-use crate::{ordered_trie_with_encoder, TrieProvider};
+use crate::{ordered_trie_with_encoder, TrieNode, TrieProvider};
 use alloc::{collections::BTreeMap, vec::Vec};
 use alloy_consensus::{Receipt, ReceiptEnvelope, ReceiptWithBloom, TxEnvelope, TxType};
 use alloy_primitives::{keccak256, Bytes, Log, B256};
 use alloy_provider::{network::eip2718::Encodable2718, Provider, ProviderBuilder};
+use alloy_rlp::Decodable;
 use alloy_rpc_types::{BlockTransactions, BlockTransactionsKind};
 use anyhow::{anyhow, Result};
 use reqwest::Url;
@@ -138,8 +139,16 @@ impl TrieNodeProvider {
 impl TrieProvider for TrieNodeProvider {
     type Error = anyhow::Error;
 
-    fn trie_node_preimage(&self, key: B256) -> Result<Bytes> {
-        self.preimages.get(&key).cloned().ok_or_else(|| anyhow!("Key not found"))
+    fn trie_node_by_hash(&self, key: B256) -> Result<TrieNode> {
+        TrieNode::decode(
+            &mut self
+                .preimages
+                .get(&key)
+                .cloned()
+                .ok_or_else(|| anyhow!("Key not found"))?
+                .as_ref(),
+        )
+        .map_err(Into::into)
     }
 
     fn bytecode_by_hash(&self, hash: B256) -> Result<Bytes> {

diff --git a/crates/mpt/src/traits.rs b/crates/mpt/src/traits.rs
@@ -1,6 +1,7 @@
 //! Contains the [TrieProvider] trait for fetching trie node preimages, contract bytecode, and
 //! headers.
 
+use crate::TrieNode;
 use alloc::string::ToString;
 use alloy_consensus::Header;
 use alloy_primitives::{Address, Bytes, B256, U256};
@@ -18,11 +19,11 @@ pub trait TrieProvider {
     /// - `key`: The key of the trie node to fetch.
     ///
     /// ## Returns
-    /// - Ok(Bytes): The trie node preimage.
+    /// - Ok(TrieNode): The trie node preimage.
     /// - Err(anyhow::Error): If the trie node preimage could not be fetched.
     ///
     /// [TrieDB]: crate::TrieDB
-    fn trie_node_preimage(&self, key: B256) -> Result<Bytes, Self::Error>;
+    fn trie_node_by_hash(&self, key: B256) -> Result<TrieNode, Self::Error>;
 
     /// Fetches the preimage of the bytecode hash provided.
     ///