From 928bcaaaab05a20b6fded6f1660591477f75c470 Mon Sep 17 00:00:00 2001
From: frisitano <francesco.risitano95@gmail.com>
Date: Thu, 17 Oct 2024 23:04:49 +0800
Subject: [PATCH] feat: introduce StateCommitment type

---
 Cargo.lock                                    |  5 ++
 crates/ethereum/node/Cargo.toml               |  2 +
 crates/ethereum/node/src/node.rs              |  2 +
 crates/exex/test-utils/Cargo.toml             |  1 +
 crates/exex/test-utils/src/lib.rs             |  1 +
 crates/node/builder/src/node.rs               |  2 +
 crates/node/types/Cargo.toml                  |  1 +
 crates/node/types/src/lib.rs                  | 49 +++++++++++++------
 crates/optimism/node/Cargo.toml               |  2 +
 crates/optimism/node/src/node.rs              |  2 +
 .../storage/provider/src/test_utils/mock.rs   |  2 +
 crates/storage/provider/src/test_utils/mod.rs |  1 +
 crates/trie/common/src/key.rs                 | 18 +++++++
 crates/trie/common/src/lib.rs                 |  3 ++
 crates/trie/db/src/commitment.rs              | 39 +++++++++++++++
 crates/trie/db/src/lib.rs                     |  2 +
 examples/custom-engine-types/Cargo.toml       |  1 +
 examples/custom-engine-types/src/main.rs      |  2 +
 18 files changed, 120 insertions(+), 15 deletions(-)
 create mode 100644 crates/trie/common/src/key.rs
 create mode 100644 crates/trie/db/src/commitment.rs

diff --git a/Cargo.lock b/Cargo.lock
index 5b62244576ef..8f5529a1f691 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2856,6 +2856,7 @@ dependencies = [
  "reth-payload-builder",
  "reth-primitives",
  "reth-tracing",
+ "reth-trie-db",
  "serde",
  "thiserror",
  "tokio",
@@ -7563,6 +7564,7 @@ dependencies = [
  "reth-provider",
  "reth-tasks",
  "reth-transaction-pool",
+ "reth-trie-db",
  "tempfile",
  "thiserror",
  "tokio",
@@ -8023,6 +8025,7 @@ dependencies = [
  "reth-tasks",
  "reth-tracing",
  "reth-transaction-pool",
+ "reth-trie-db",
  "revm",
  "serde_json",
  "tokio",
@@ -8085,6 +8088,7 @@ dependencies = [
  "reth-engine-primitives",
  "reth-primitives",
  "reth-primitives-traits",
+ "reth-trie-db",
 ]
 
 [[package]]
@@ -8245,6 +8249,7 @@ dependencies = [
  "reth-revm",
  "reth-tracing",
  "reth-transaction-pool",
+ "reth-trie-db",
  "revm",
  "serde",
  "serde_json",
diff --git a/crates/ethereum/node/Cargo.toml b/crates/ethereum/node/Cargo.toml
index 98224b99e26b..83b62034740b 100644
--- a/crates/ethereum/node/Cargo.toml
+++ b/crates/ethereum/node/Cargo.toml
@@ -31,6 +31,7 @@ reth-node-api.workspace = true
 reth-chainspec.workspace = true
 reth-primitives.workspace = true
 reth-revm = { workspace = true, features = ["std"] }
+reth-trie-db.workspace = true
 
 # revm with required ethereum features
 revm = { workspace = true, features = ["secp256k1", "blst", "c-kzg"] }
@@ -72,6 +73,7 @@ test-utils = [
 	"reth-db/test-utils",
 	"reth-provider/test-utils",
 	"reth-transaction-pool/test-utils",
+	"reth-trie-db/test-utils",
 	"revm/test-utils",
 	"reth-evm/test-utils"
 ]
diff --git a/crates/ethereum/node/src/node.rs b/crates/ethereum/node/src/node.rs
index dbd6ce0a134a..1942f8a9e56b 100644
--- a/crates/ethereum/node/src/node.rs
+++ b/crates/ethereum/node/src/node.rs
@@ -34,6 +34,7 @@ use reth_transaction_pool::{
     blobstore::DiskFileBlobStore, EthTransactionPool, TransactionPool,
     TransactionValidationTaskExecutor,
 };
+use reth_trie_db::MerklePatriciaTrie;
 
 use crate::{EthEngineTypes, EthEvmConfig};
 
@@ -81,6 +82,7 @@ impl EthereumNode {
 impl NodeTypes for EthereumNode {
     type Primitives = EthPrimitives;
     type ChainSpec = ChainSpec;
+    type StateCommitment = MerklePatriciaTrie;
 }
 
 impl NodeTypesWithEngine for EthereumNode {
diff --git a/crates/exex/test-utils/Cargo.toml b/crates/exex/test-utils/Cargo.toml
index b850295a332b..cd0e0831b496 100644
--- a/crates/exex/test-utils/Cargo.toml
+++ b/crates/exex/test-utils/Cargo.toml
@@ -31,6 +31,7 @@ reth-primitives.workspace = true
 reth-provider = { workspace = true, features = ["test-utils"] }
 reth-tasks.workspace = true
 reth-transaction-pool = { workspace = true, features = ["test-utils"] }
+reth-trie-db.workspace = true
 
 ## async
 futures-util.workspace = true
diff --git a/crates/exex/test-utils/src/lib.rs b/crates/exex/test-utils/src/lib.rs
index 1f6ea75ce6d9..06aa8c81c7cd 100644
--- a/crates/exex/test-utils/src/lib.rs
+++ b/crates/exex/test-utils/src/lib.rs
@@ -119,6 +119,7 @@ pub struct TestNode;
 impl NodeTypes for TestNode {
     type Primitives = ();
     type ChainSpec = ChainSpec;
+    type StateCommitment = reth_trie_db::MerklePatriciaTrie;
 }
 
 impl NodeTypesWithEngine for TestNode {
diff --git a/crates/node/builder/src/node.rs b/crates/node/builder/src/node.rs
index 3b2f467d61c0..62c710ea8022 100644
--- a/crates/node/builder/src/node.rs
+++ b/crates/node/builder/src/node.rs
@@ -69,6 +69,8 @@ where
     type Primitives = <N::Types as NodeTypes>::Primitives;
 
     type ChainSpec = <N::Types as NodeTypes>::ChainSpec;
+
+    type StateCommitment = <N::Types as NodeTypes>::StateCommitment;
 }
 
 impl<N, C, AO> NodeTypesWithEngine for AnyNode<N, C, AO>
diff --git a/crates/node/types/Cargo.toml b/crates/node/types/Cargo.toml
index 5747abe9c34b..21facae5460a 100644
--- a/crates/node/types/Cargo.toml
+++ b/crates/node/types/Cargo.toml
@@ -17,3 +17,4 @@ reth-db-api.workspace = true
 reth-engine-primitives.workspace = true
 reth-primitives.workspace = true
 reth-primitives-traits.workspace = true
+reth-trie-db.workspace = true
diff --git a/crates/node/types/src/lib.rs b/crates/node/types/src/lib.rs
index 5ba03e6795ab..38e194bd4fbf 100644
--- a/crates/node/types/src/lib.rs
+++ b/crates/node/types/src/lib.rs
@@ -18,6 +18,7 @@ use reth_db_api::{
     Database,
 };
 use reth_engine_primitives::EngineTypes;
+use reth_trie_db::StateCommitment;
 
 /// Configures all the primitive types of the node.
 pub trait NodePrimitives {
@@ -39,6 +40,8 @@ pub trait NodeTypes: Send + Sync + Unpin + 'static {
     type Primitives: NodePrimitives;
     /// The type used for configuration of the EVM.
     type ChainSpec: EthChainSpec;
+    /// The type used to perform state commitment operations.
+    type StateCommitment: StateCommitment;
 }
 
 /// The type that configures an Ethereum-like node with an engine for consensus.
@@ -89,6 +92,7 @@ where
 {
     type Primitives = Types::Primitives;
     type ChainSpec = Types::ChainSpec;
+    type StateCommitment = Types::StateCommitment;
 }
 
 impl<Types, DB> NodeTypesWithEngine for NodeTypesWithDBAdapter<Types, DB>
@@ -109,70 +113,85 @@ where
 
 /// A [`NodeTypes`] type builder.
 #[derive(Default, Debug)]
-pub struct AnyNodeTypes<P = (), C = ()>(PhantomData<P>, PhantomData<C>);
+pub struct AnyNodeTypes<P = (), C = (), S = ()>(PhantomData<P>, PhantomData<C>, PhantomData<S>);
 
-impl<P, C> AnyNodeTypes<P, C> {
+impl<P, C, S> AnyNodeTypes<P, C, S> {
     /// Sets the `Primitives` associated type.
-    pub const fn primitives<T>(self) -> AnyNodeTypes<T, C> {
-        AnyNodeTypes::<T, C>(PhantomData::<T>, PhantomData::<C>)
+    pub const fn primitives<T>(self) -> AnyNodeTypes<T, C, S> {
+        AnyNodeTypes::<T, C, S>(PhantomData::<T>, PhantomData::<C>, PhantomData::<S>)
     }
 
     /// Sets the `ChainSpec` associated type.
-    pub const fn chain_spec<T>(self) -> AnyNodeTypes<P, T> {
-        AnyNodeTypes::<P, T>(PhantomData::<P>, PhantomData::<T>)
+    pub const fn chain_spec<T>(self) -> AnyNodeTypes<P, T, S> {
+        AnyNodeTypes::<P, T, S>(PhantomData::<P>, PhantomData::<T>, PhantomData::<S>)
+    }
+
+    /// Sets the `StateCommitment` associated type.
+    pub const fn state_commitment<T>(self) -> AnyNodeTypes<P, C, T> {
+        AnyNodeTypes::<P, C, T>(PhantomData::<P>, PhantomData::<C>, PhantomData::<T>)
     }
 }
 
-impl<P, C> NodeTypes for AnyNodeTypes<P, C>
+impl<P, C, S> NodeTypes for AnyNodeTypes<P, C, S>
 where
     P: NodePrimitives + Send + Sync + Unpin + 'static,
     C: EthChainSpec + 'static,
+    S: StateCommitment,
 {
     type Primitives = P;
     type ChainSpec = C;
+    type StateCommitment = S;
 }
 
 /// A [`NodeTypesWithEngine`] type builder.
 #[derive(Default, Debug)]
-pub struct AnyNodeTypesWithEngine<P = (), E = (), C = ()> {
+pub struct AnyNodeTypesWithEngine<P = (), E = (), C = (), S = ()> {
     /// Embedding the basic node types.
-    base: AnyNodeTypes<P, C>,
+    base: AnyNodeTypes<P, C, S>,
     /// Phantom data for the engine.
     _engine: PhantomData<E>,
 }
 
-impl<P, E, C> AnyNodeTypesWithEngine<P, E, C> {
+impl<P, E, C, S> AnyNodeTypesWithEngine<P, E, C, S> {
     /// Sets the `Primitives` associated type.
-    pub const fn primitives<T>(self) -> AnyNodeTypesWithEngine<T, E, C> {
+    pub const fn primitives<T>(self) -> AnyNodeTypesWithEngine<T, E, C, S> {
         AnyNodeTypesWithEngine { base: self.base.primitives::<T>(), _engine: PhantomData }
     }
 
     /// Sets the `Engine` associated type.
-    pub const fn engine<T>(self) -> AnyNodeTypesWithEngine<P, T, C> {
+    pub const fn engine<T>(self) -> AnyNodeTypesWithEngine<P, T, C, S> {
         AnyNodeTypesWithEngine { base: self.base, _engine: PhantomData::<T> }
     }
 
     /// Sets the `ChainSpec` associated type.
-    pub const fn chain_spec<T>(self) -> AnyNodeTypesWithEngine<P, E, T> {
+    pub const fn chain_spec<T>(self) -> AnyNodeTypesWithEngine<P, E, T, S> {
         AnyNodeTypesWithEngine { base: self.base.chain_spec::<T>(), _engine: PhantomData }
     }
+
+    /// Sets the `StateCommitment` associated type.
+    pub const fn state_commitment<T>(self) -> AnyNodeTypesWithEngine<P, E, C, T> {
+        AnyNodeTypesWithEngine { base: self.base.state_commitment::<T>(), _engine: PhantomData }
+    }
 }
 
-impl<P, E, C> NodeTypes for AnyNodeTypesWithEngine<P, E, C>
+impl<P, E, C, S> NodeTypes for AnyNodeTypesWithEngine<P, E, C, S>
 where
     P: NodePrimitives + Send + Sync + Unpin + 'static,
     E: EngineTypes + Send + Sync + Unpin,
     C: EthChainSpec + 'static,
+    S: StateCommitment,
 {
     type Primitives = P;
     type ChainSpec = C;
+    type StateCommitment = S;
 }
 
-impl<P, E, C> NodeTypesWithEngine for AnyNodeTypesWithEngine<P, E, C>
+impl<P, E, C, S> NodeTypesWithEngine for AnyNodeTypesWithEngine<P, E, C, S>
 where
     P: NodePrimitives + Send + Sync + Unpin + 'static,
     E: EngineTypes + Send + Sync + Unpin,
     C: EthChainSpec + 'static,
+    S: StateCommitment,
 {
     type Engine = E;
 }
diff --git a/crates/optimism/node/Cargo.toml b/crates/optimism/node/Cargo.toml
index 37cf4a328ea0..deabbac52499 100644
--- a/crates/optimism/node/Cargo.toml
+++ b/crates/optimism/node/Cargo.toml
@@ -28,6 +28,7 @@ reth-network.workspace = true
 reth-evm.workspace = true
 reth-revm = { workspace = true, features = ["std"] }
 reth-beacon-consensus.workspace = true
+reth-trie-db.workspace = true
 
 # op-reth
 reth-optimism-payload-builder.workspace = true
@@ -99,5 +100,6 @@ test-utils = [
 	"reth-db/test-utils",
 	"reth-provider/test-utils",
 	"reth-transaction-pool/test-utils",
+	"reth-trie-db/test-utils",
 	"revm/test-utils"
 ]
diff --git a/crates/optimism/node/src/node.rs b/crates/optimism/node/src/node.rs
index 9492bb8c429a..4328a55fb15d 100644
--- a/crates/optimism/node/src/node.rs
+++ b/crates/optimism/node/src/node.rs
@@ -30,6 +30,7 @@ use reth_transaction_pool::{
     blobstore::DiskFileBlobStore, CoinbaseTipOrdering, TransactionPool,
     TransactionValidationTaskExecutor,
 };
+use reth_trie_db::MerklePatriciaTrie;
 
 use crate::{
     args::RollupArgs,
@@ -122,6 +123,7 @@ where
 impl NodeTypes for OptimismNode {
     type Primitives = OpPrimitives;
     type ChainSpec = OpChainSpec;
+    type StateCommitment = MerklePatriciaTrie;
 }
 
 impl NodeTypesWithEngine for OptimismNode {
diff --git a/crates/storage/provider/src/test_utils/mock.rs b/crates/storage/provider/src/test_utils/mock.rs
index 1053b4778fdd..e2e08e61a86f 100644
--- a/crates/storage/provider/src/test_utils/mock.rs
+++ b/crates/storage/provider/src/test_utils/mock.rs
@@ -34,6 +34,7 @@ use reth_trie::{
     updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage, MultiProof, StorageProof,
     TrieInput,
 };
+use reth_trie_db::MerklePatriciaTrie;
 use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
 use std::{
     collections::BTreeMap,
@@ -157,6 +158,7 @@ pub struct MockNode;
 impl NodeTypes for MockNode {
     type Primitives = ();
     type ChainSpec = ChainSpec;
+    type StateCommitment = MerklePatriciaTrie;
 }
 
 impl DatabaseProviderFactory for MockEthProvider {
diff --git a/crates/storage/provider/src/test_utils/mod.rs b/crates/storage/provider/src/test_utils/mod.rs
index 2200781096d0..c0e80930b318 100644
--- a/crates/storage/provider/src/test_utils/mod.rs
+++ b/crates/storage/provider/src/test_utils/mod.rs
@@ -25,6 +25,7 @@ pub type MockNodeTypes = reth_node_types::AnyNodeTypesWithEngine<
     (),
     reth_ethereum_engine_primitives::EthEngineTypes,
     reth_chainspec::ChainSpec,
+    reth_trie_db::MerklePatriciaTrie,
 >;
 
 /// Mock [`reth_node_types::NodeTypesWithDB`] for testing.
diff --git a/crates/trie/common/src/key.rs b/crates/trie/common/src/key.rs
new file mode 100644
index 000000000000..9e440d199fa1
--- /dev/null
+++ b/crates/trie/common/src/key.rs
@@ -0,0 +1,18 @@
+use alloy_primitives::B256;
+use revm_primitives::keccak256;
+
+/// Trait for hashing keys in state.
+pub trait KeyHasher: Default + Clone + Send + Sync + 'static {
+    /// Hashes the given bytes into a 256-bit hash.
+    fn hash_key<T: AsRef<[u8]>>(bytes: T) -> B256;
+}
+
+/// A key hasher that uses the Keccak-256 hash function.
+#[derive(Clone, Debug, Default)]
+pub struct KeccakKeyHasher;
+
+impl KeyHasher for KeccakKeyHasher {
+    fn hash_key<T: AsRef<[u8]>>(bytes: T) -> B256 {
+        keccak256(bytes)
+    }
+}
diff --git a/crates/trie/common/src/lib.rs b/crates/trie/common/src/lib.rs
index bdec36028b94..7645ebd3a1cb 100644
--- a/crates/trie/common/src/lib.rs
+++ b/crates/trie/common/src/lib.rs
@@ -14,6 +14,9 @@ pub mod hash_builder;
 mod account;
 pub use account::TrieAccount;
 
+mod key;
+pub use key::{KeccakKeyHasher, KeyHasher};
+
 mod nibbles;
 pub use nibbles::{Nibbles, StoredNibbles, StoredNibblesSubKey};
 
diff --git a/crates/trie/db/src/commitment.rs b/crates/trie/db/src/commitment.rs
new file mode 100644
index 000000000000..c608aefff8ae
--- /dev/null
+++ b/crates/trie/db/src/commitment.rs
@@ -0,0 +1,39 @@
+use crate::{
+    DatabaseHashedCursorFactory, DatabaseProof, DatabaseStateRoot, DatabaseStorageRoot,
+    DatabaseTrieCursorFactory, DatabaseTrieWitness,
+};
+use reth_db::transaction::DbTx;
+use reth_trie::{
+    proof::Proof, witness::TrieWitness, KeccakKeyHasher, KeyHasher, StateRoot, StorageRoot,
+};
+
+/// The `StateCommitment` trait provides associated types for state commitment operations.
+pub trait StateCommitment: std::fmt::Debug + Send + Sync + Unpin + 'static {
+    /// The state root type.
+    type StateRoot<'a, TX: DbTx + 'a>: DatabaseStateRoot<'a, TX>;
+    /// The storage root type.
+    type StorageRoot<'a, TX: DbTx + 'a>: DatabaseStorageRoot<'a, TX>;
+    /// The state proof type.
+    type StateProof<'a, TX: DbTx + 'a>: DatabaseProof<'a, TX>;
+    /// The state witness type.
+    type StateWitness<'a, TX: DbTx + 'a>: DatabaseTrieWitness<'a, TX>;
+    /// The key hasher type.
+    type KeyHasher: KeyHasher;
+}
+
+/// The state commitment type for Ethereum's Merkle Patricia Trie.
+#[derive(Debug)]
+#[non_exhaustive]
+pub struct MerklePatriciaTrie;
+
+impl StateCommitment for MerklePatriciaTrie {
+    type StateRoot<'a, TX: DbTx + 'a> =
+        StateRoot<DatabaseTrieCursorFactory<'a, TX>, DatabaseHashedCursorFactory<'a, TX>>;
+    type StorageRoot<'a, TX: DbTx + 'a> =
+        StorageRoot<DatabaseTrieCursorFactory<'a, TX>, DatabaseHashedCursorFactory<'a, TX>>;
+    type StateProof<'a, TX: DbTx + 'a> =
+        Proof<DatabaseTrieCursorFactory<'a, TX>, DatabaseHashedCursorFactory<'a, TX>>;
+    type StateWitness<'a, TX: DbTx + 'a> =
+        TrieWitness<DatabaseTrieCursorFactory<'a, TX>, DatabaseHashedCursorFactory<'a, TX>>;
+    type KeyHasher = KeccakKeyHasher;
+}
diff --git a/crates/trie/db/src/lib.rs b/crates/trie/db/src/lib.rs
index 3a9b1e328239..27c18af6cbfd 100644
--- a/crates/trie/db/src/lib.rs
+++ b/crates/trie/db/src/lib.rs
@@ -1,5 +1,6 @@
 //! An integration of [`reth-trie`] with [`reth-db`].
 
+mod commitment;
 mod hashed_cursor;
 mod prefix_set;
 mod proof;
@@ -8,6 +9,7 @@ mod storage;
 mod trie_cursor;
 mod witness;
 
+pub use commitment::{MerklePatriciaTrie, StateCommitment};
 pub use hashed_cursor::{
     DatabaseHashedAccountCursor, DatabaseHashedCursorFactory, DatabaseHashedStorageCursor,
 };
diff --git a/examples/custom-engine-types/Cargo.toml b/examples/custom-engine-types/Cargo.toml
index f826451d2038..1fbb3c4947ac 100644
--- a/examples/custom-engine-types/Cargo.toml
+++ b/examples/custom-engine-types/Cargo.toml
@@ -16,6 +16,7 @@ reth-basic-payload-builder.workspace = true
 reth-ethereum-payload-builder.workspace = true
 reth-node-ethereum = { workspace = true, features = ["test-utils"] }
 reth-tracing.workspace = true
+reth-trie-db.workspace = true
 alloy-genesis.workspace = true
 alloy-rpc-types = { workspace = true, features = ["engine"] }
 alloy-primitives.workspace = true
diff --git a/examples/custom-engine-types/src/main.rs b/examples/custom-engine-types/src/main.rs
index 30f89a0b9d78..a48c44b200b9 100644
--- a/examples/custom-engine-types/src/main.rs
+++ b/examples/custom-engine-types/src/main.rs
@@ -70,6 +70,7 @@ use reth_payload_builder::{
 };
 use reth_primitives::Withdrawals;
 use reth_tracing::{RethTracer, Tracer};
+use reth_trie_db::MerklePatriciaTrie;
 
 /// A custom payload attributes type.
 #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
@@ -228,6 +229,7 @@ struct MyCustomNode;
 impl NodeTypes for MyCustomNode {
     type Primitives = ();
     type ChainSpec = ChainSpec;
+    type StateCommitment = MerklePatriciaTrie;
 }
 
 /// Configure the node types with the custom engine types