From e4c6294fc5531513060887d822c5cc591d48ba8d 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               |  1 +
 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                  |  3 +-
 crates/node/types/src/lib.rs                  | 49 +++++++++++++------
 crates/optimism/node/Cargo.toml               |  1 +
 crates/optimism/node/src/node.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              | 38 ++++++++++++++
 crates/trie/db/src/lib.rs                     |  2 +
 examples/custom-engine-types/Cargo.toml       |  1 +
 examples/custom-engine-types/src/main.rs      |  2 +
 17 files changed, 116 insertions(+), 16 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 e434e80cfde2..fc07aed94ca2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2754,6 +2754,7 @@ dependencies = [
  "reth-payload-builder",
  "reth-primitives",
  "reth-tracing",
+ "reth-trie-db",
  "serde",
  "thiserror",
  "tokio",
@@ -7504,6 +7505,7 @@ dependencies = [
  "reth-provider",
  "reth-tasks",
  "reth-transaction-pool",
+ "reth-trie-db",
  "tempfile",
  "thiserror",
  "tokio",
@@ -7960,6 +7962,7 @@ dependencies = [
  "reth-tasks",
  "reth-tracing",
  "reth-transaction-pool",
+ "reth-trie-db",
  "revm",
  "serde_json",
  "tokio",
@@ -8020,6 +8023,7 @@ dependencies = [
  "reth-chainspec",
  "reth-db-api",
  "reth-engine-primitives",
+ "reth-trie-db",
 ]
 
 [[package]]
@@ -8185,6 +8189,7 @@ dependencies = [
  "reth-rpc-types-compat",
  "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 7a323f91d877..d264f3daea56 100644
--- a/crates/ethereum/node/Cargo.toml
+++ b/crates/ethereum/node/Cargo.toml
@@ -30,6 +30,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"] }
diff --git a/crates/ethereum/node/src/node.rs b/crates/ethereum/node/src/node.rs
index 82f313fbb0b2..27dca26c022c 100644
--- a/crates/ethereum/node/src/node.rs
+++ b/crates/ethereum/node/src/node.rs
@@ -30,6 +30,7 @@ use reth_transaction_pool::{
     blobstore::DiskFileBlobStore, EthTransactionPool, TransactionPool,
     TransactionValidationTaskExecutor,
 };
+use reth_trie_db::MerklePatriciaTrie;
 
 use crate::{EthEngineTypes, EthEvmConfig};
 
@@ -71,6 +72,7 @@ impl EthereumNode {
 impl NodeTypes for EthereumNode {
     type Primitives = ();
     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 8488cdb8b731..a88d0832632a 100644
--- a/crates/exex/test-utils/Cargo.toml
+++ b/crates/exex/test-utils/Cargo.toml
@@ -32,6 +32,7 @@ reth-provider = { workspace = true, features = ["test-utils"] }
 reth-tasks.workspace = true
 reth-transaction-pool = { workspace = true, features = ["test-utils"] }
 reth-ethereum-engine-primitives.workspace = true
+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 9e17013c4a5e..07db668aec52 100644
--- a/crates/exex/test-utils/src/lib.rs
+++ b/crates/exex/test-utils/src/lib.rs
@@ -123,6 +123,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 3e3d5b696c39..b04b763125ba 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 f04925d9cd41..9e99287d844d 100644
--- a/crates/node/types/Cargo.toml
+++ b/crates/node/types/Cargo.toml
@@ -14,4 +14,5 @@ workspace = true
 # reth
 reth-chainspec.workspace = true
 reth-db-api.workspace = true
-reth-engine-primitives.workspace = true
\ No newline at end of file
+reth-engine-primitives.workspace = true
+reth-trie-db.workspace = true
diff --git a/crates/node/types/src/lib.rs b/crates/node/types/src/lib.rs
index 2c72e02d3edc..36cb68ac8c65 100644
--- a/crates/node/types/src/lib.rs
+++ b/crates/node/types/src/lib.rs
@@ -16,6 +16,7 @@ use reth_db_api::{
     Database,
 };
 use reth_engine_primitives::EngineTypes;
+use reth_trie_db::StateCommitment;
 
 /// Configures all the primitive types of the node.
 // TODO(mattsse): this is currently a placeholder
@@ -34,6 +35,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.
@@ -84,6 +87,7 @@ where
 {
     type Primitives = Types::Primitives;
     type ChainSpec = Types::ChainSpec;
+    type StateCommitment = Types::StateCommitment;
 }
 
 impl<Types, DB> NodeTypesWithEngine for NodeTypesWithDBAdapter<Types, DB>
@@ -104,70 +108,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<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<T>(self) -> AnyNodeTypesWithEngine<P, E, C, T> {
+        AnyNodeTypesWithEngine { base: self.base.state::<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 8e359e602657..d1338e4c4591 100644
--- a/crates/optimism/node/Cargo.toml
+++ b/crates/optimism/node/Cargo.toml
@@ -33,6 +33,7 @@ reth-discv5.workspace = true
 reth-rpc-eth-types.workspace = true
 reth-rpc-eth-api.workspace = true
 reth-rpc.workspace = true
+reth-trie-db.workspace = true
 
 # op-reth
 reth-optimism-payload-builder.workspace = true
diff --git a/crates/optimism/node/src/node.rs b/crates/optimism/node/src/node.rs
index 648da85d0bb4..7f61e0b71df5 100644
--- a/crates/optimism/node/src/node.rs
+++ b/crates/optimism/node/src/node.rs
@@ -28,6 +28,7 @@ use reth_transaction_pool::{
     blobstore::DiskFileBlobStore, CoinbaseTipOrdering, TransactionPool,
     TransactionValidationTaskExecutor,
 };
+use reth_trie_db::MerklePatriciaTrie;
 
 use crate::{
     args::RollupArgs,
@@ -115,6 +116,7 @@ where
 impl NodeTypes for OptimismNode {
     type Primitives = ();
     type ChainSpec = OpChainSpec;
+    type StateCommitment = MerklePatriciaTrie;
 }
 
 impl NodeTypesWithEngine for OptimismNode {
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..30b1d07469af
--- /dev/null
+++ b/crates/trie/db/src/commitment.rs
@@ -0,0 +1,38 @@
+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)]
+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 f833da86236e..ab0ec5f52dbe 100644
--- a/examples/custom-engine-types/src/main.rs
+++ b/examples/custom-engine-types/src/main.rs
@@ -67,6 +67,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)]
@@ -219,6 +220,7 @@ struct MyCustomNode;
 impl NodeTypes for MyCustomNode {
     type Primitives = ();
     type ChainSpec = ChainSpec;
+    type StateCommitment = MerklePatriciaTrie;
 }
 
 /// Configure the node types with the custom engine types