From 83af493179393c81dc2c48374306f2b88cc075db Mon Sep 17 00:00:00 2001
From: frisitano <35734660+frisitano@users.noreply.github.com>
Date: Wed, 27 Nov 2024 00:48:32 +0800
Subject: [PATCH] Introduce StateCommitment in StateProviders (#12602)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>
---
 Cargo.lock                                    |  1 +
 crates/stages/stages/src/stages/execution.rs  |  7 ++--
 .../provider/src/providers/consistent.rs      |  2 +-
 .../src/providers/database/provider.rs        | 17 +++++----
 .../src/providers/state/historical.rs         | 35 ++++++++++++-------
 .../provider/src/providers/state/latest.rs    | 24 +++++++++----
 crates/storage/storage-api/Cargo.toml         |  1 +
 crates/storage/storage-api/src/state.rs       |  7 ++++
 8 files changed, 65 insertions(+), 29 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 976058b5d4b2..2b754e8af291 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -9230,6 +9230,7 @@ dependencies = [
  "reth-stages-types",
  "reth-storage-errors",
  "reth-trie",
+ "reth-trie-db",
 ]
 
 [[package]]
diff --git a/crates/stages/stages/src/stages/execution.rs b/crates/stages/stages/src/stages/execution.rs
index e1c25c7d5fa6..f7ee52fbfe34 100644
--- a/crates/stages/stages/src/stages/execution.rs
+++ b/crates/stages/stages/src/stages/execution.rs
@@ -16,8 +16,8 @@ use reth_primitives_traits::{format_gas_throughput, Block, BlockBody, NodePrimit
 use reth_provider::{
     providers::{StaticFileProvider, StaticFileProviderRWRefMut, StaticFileWriter},
     BlockHashReader, BlockReader, DBProvider, HeaderProvider, LatestStateProviderRef,
-    OriginalValuesKnown, ProviderError, StateChangeWriter, StateWriter, StaticFileProviderFactory,
-    StatsReader, StorageLocation, TransactionVariant,
+    OriginalValuesKnown, ProviderError, StateChangeWriter, StateCommitmentProvider, StateWriter,
+    StaticFileProviderFactory, StatsReader, StorageLocation, TransactionVariant,
 };
 use reth_prune_types::PruneModes;
 use reth_revm::database::StateProviderDatabase;
@@ -180,7 +180,8 @@ where
         + StatsReader
         + StateChangeWriter
         + BlockHashReader
-        + StateWriter,
+        + StateWriter
+        + StateCommitmentProvider,
 {
     /// Return the id of the stage
     fn id(&self) -> StageId {
diff --git a/crates/storage/provider/src/providers/consistent.rs b/crates/storage/provider/src/providers/consistent.rs
index 86921a7f56e7..64a940190db0 100644
--- a/crates/storage/provider/src/providers/consistent.rs
+++ b/crates/storage/provider/src/providers/consistent.rs
@@ -115,7 +115,7 @@ impl<N: ProviderNodeTypes> ConsistentProvider<N> {
             Ok(self.block_state_provider_ref(state)?.boxed())
         } else {
             trace!(target: "providers::blockchain", "Using database state for latest state provider");
-            self.storage_provider.latest()
+            Ok(self.storage_provider.latest())
         }
     }
 
diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs
index e50f9760d6ba..750a186504a0 100644
--- a/crates/storage/provider/src/providers/database/provider.rs
+++ b/crates/storage/provider/src/providers/database/provider.rs
@@ -15,9 +15,10 @@ use crate::{
     HeaderSyncGapProvider, HistoricalStateProvider, HistoricalStateProviderRef, HistoryWriter,
     LatestStateProvider, LatestStateProviderRef, OriginalValuesKnown, ProviderError,
     PruneCheckpointReader, PruneCheckpointWriter, RevertsInit, StageCheckpointReader,
-    StateChangeWriter, StateProviderBox, StateReader, StateWriter, StaticFileProviderFactory,
-    StatsReader, StorageLocation, StorageReader, StorageTrieWriter, TransactionVariant,
-    TransactionsProvider, TransactionsProviderExt, TrieWriter, WithdrawalsProvider,
+    StateChangeWriter, StateCommitmentProvider, StateProviderBox, StateReader, StateWriter,
+    StaticFileProviderFactory, StatsReader, StorageLocation, StorageReader, StorageTrieWriter,
+    TransactionVariant, TransactionsProvider, TransactionsProviderExt, TrieWriter,
+    WithdrawalsProvider,
 };
 use alloy_consensus::Header;
 use alloy_eips::{
@@ -157,10 +158,10 @@ impl<TX, N: NodeTypes> DatabaseProvider<TX, N> {
 }
 
 impl<TX: DbTx + 'static, N: NodeTypes> DatabaseProvider<TX, N> {
-    /// State provider for latest block
-    pub fn latest<'a>(&'a self) -> ProviderResult<Box<dyn StateProvider + 'a>> {
+    /// State provider for latest state
+    pub fn latest<'a>(&'a self) -> Box<dyn StateProvider + 'a> {
         trace!(target: "providers::db", "Returning latest state provider");
-        Ok(Box::new(LatestStateProviderRef::new(self)))
+        Box::new(LatestStateProviderRef::new(self))
     }
 
     /// Storage provider for state at that given block hash
@@ -378,6 +379,10 @@ impl<TX: DbTx + 'static, N: NodeTypes> TryIntoHistoricalStateProvider for Databa
     }
 }
 
+impl<TX: DbTx + 'static, N: NodeTypes> StateCommitmentProvider for DatabaseProvider<TX, N> {
+    type StateCommitment = N::StateCommitment;
+}
+
 impl<Tx: DbTx + DbTxMut + 'static, N: NodeTypesForProvider + 'static> DatabaseProvider<Tx, N> {
     // TODO: uncomment below, once `reth debug_cmd` has been feature gated with dev.
     // #[cfg(any(test, feature = "test-utils"))]
diff --git a/crates/storage/provider/src/providers/state/historical.rs b/crates/storage/provider/src/providers/state/historical.rs
index ca036844f659..ad36a4a5ab3e 100644
--- a/crates/storage/provider/src/providers/state/historical.rs
+++ b/crates/storage/provider/src/providers/state/historical.rs
@@ -15,7 +15,9 @@ use reth_db_api::{
     transaction::DbTx,
 };
 use reth_primitives::{Account, Bytecode};
-use reth_storage_api::{BlockNumReader, DBProvider, StateProofProvider, StorageRootProvider};
+use reth_storage_api::{
+    BlockNumReader, DBProvider, StateCommitmentProvider, StateProofProvider, StorageRootProvider,
+};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
     proof::{Proof, StorageProof},
@@ -59,7 +61,9 @@ pub enum HistoryInfo {
     MaybeInPlainState,
 }
 
-impl<'b, Provider: DBProvider + BlockNumReader> HistoricalStateProviderRef<'b, Provider> {
+impl<'b, Provider: DBProvider + BlockNumReader + StateCommitmentProvider>
+    HistoricalStateProviderRef<'b, Provider>
+{
     /// Create new `StateProvider` for historical block number
     pub fn new(provider: &'b Provider, block_number: BlockNumber) -> Self {
         Self { provider, block_number, lowest_available_blocks: Default::default() }
@@ -240,7 +244,7 @@ impl<Provider: DBProvider + BlockNumReader> HistoricalStateProviderRef<'_, Provi
     }
 }
 
-impl<Provider: DBProvider + BlockNumReader> AccountReader
+impl<Provider: DBProvider + BlockNumReader + StateCommitmentProvider> AccountReader
     for HistoricalStateProviderRef<'_, Provider>
 {
     /// Get basic account information.
@@ -281,7 +285,7 @@ impl<Provider: DBProvider + BlockNumReader + BlockHashReader> BlockHashReader
     }
 }
 
-impl<Provider: DBProvider + BlockNumReader> StateRootProvider
+impl<Provider: DBProvider + BlockNumReader + StateCommitmentProvider> StateRootProvider
     for HistoricalStateProviderRef<'_, Provider>
 {
     fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
@@ -317,7 +321,7 @@ impl<Provider: DBProvider + BlockNumReader> StateRootProvider
     }
 }
 
-impl<Provider: DBProvider + BlockNumReader> StorageRootProvider
+impl<Provider: DBProvider + BlockNumReader + StateCommitmentProvider> StorageRootProvider
     for HistoricalStateProviderRef<'_, Provider>
 {
     fn storage_root(
@@ -356,7 +360,7 @@ impl<Provider: DBProvider + BlockNumReader> StorageRootProvider
     }
 }
 
-impl<Provider: DBProvider + BlockNumReader> StateProofProvider
+impl<Provider: DBProvider + BlockNumReader + StateCommitmentProvider> StateProofProvider
     for HistoricalStateProviderRef<'_, Provider>
 {
     /// Get account and storage proofs.
@@ -390,8 +394,8 @@ impl<Provider: DBProvider + BlockNumReader> StateProofProvider
     }
 }
 
-impl<Provider: DBProvider + BlockNumReader + BlockHashReader> StateProvider
-    for HistoricalStateProviderRef<'_, Provider>
+impl<Provider: DBProvider + BlockNumReader + BlockHashReader + StateCommitmentProvider>
+    StateProvider for HistoricalStateProviderRef<'_, Provider>
 {
     /// Get storage.
     fn storage(
@@ -441,7 +445,9 @@ pub struct HistoricalStateProvider<Provider> {
     lowest_available_blocks: LowestAvailableBlocks,
 }
 
-impl<Provider: DBProvider + BlockNumReader> HistoricalStateProvider<Provider> {
+impl<Provider: DBProvider + BlockNumReader + StateCommitmentProvider>
+    HistoricalStateProvider<Provider>
+{
     /// Create new `StateProvider` for historical block number
     pub fn new(provider: Provider, block_number: BlockNumber) -> Self {
         Self { provider, block_number, lowest_available_blocks: Default::default() }
@@ -477,7 +483,7 @@ impl<Provider: DBProvider + BlockNumReader> HistoricalStateProvider<Provider> {
 }
 
 // Delegates all provider impls to [HistoricalStateProviderRef]
-delegate_provider_impls!(HistoricalStateProvider<Provider> where [Provider: DBProvider + BlockNumReader + BlockHashReader]);
+delegate_provider_impls!(HistoricalStateProvider<Provider> where [Provider: DBProvider + BlockNumReader + BlockHashReader + StateCommitmentProvider]);
 
 /// Lowest blocks at which different parts of the state are available.
 /// They may be [Some] if pruning is enabled.
@@ -521,7 +527,10 @@ mod tests {
         transaction::{DbTx, DbTxMut},
     };
     use reth_primitives::{Account, StorageEntry};
-    use reth_storage_api::{BlockHashReader, BlockNumReader, DBProvider, DatabaseProviderFactory};
+    use reth_storage_api::{
+        BlockHashReader, BlockNumReader, DBProvider, DatabaseProviderFactory,
+        StateCommitmentProvider,
+    };
     use reth_storage_errors::provider::ProviderError;
 
     const ADDRESS: Address = address!("0000000000000000000000000000000000000001");
@@ -530,7 +539,9 @@ mod tests {
 
     const fn assert_state_provider<T: StateProvider>() {}
     #[allow(dead_code)]
-    const fn assert_historical_state_provider<T: DBProvider + BlockNumReader + BlockHashReader>() {
+    const fn assert_historical_state_provider<
+        T: DBProvider + BlockNumReader + BlockHashReader + StateCommitmentProvider,
+    >() {
         assert_state_provider::<HistoricalStateProvider<T>>();
     }
 
diff --git a/crates/storage/provider/src/providers/state/latest.rs b/crates/storage/provider/src/providers/state/latest.rs
index 67dd1e744711..a2ec4972d105 100644
--- a/crates/storage/provider/src/providers/state/latest.rs
+++ b/crates/storage/provider/src/providers/state/latest.rs
@@ -9,7 +9,9 @@ use alloy_primitives::{
 use reth_db::tables;
 use reth_db_api::{cursor::DbDupCursorRO, transaction::DbTx};
 use reth_primitives::{Account, Bytecode};
-use reth_storage_api::{DBProvider, StateProofProvider, StorageRootProvider};
+use reth_storage_api::{
+    DBProvider, StateCommitmentProvider, StateProofProvider, StorageRootProvider,
+};
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{
     proof::{Proof, StorageProof},
@@ -62,7 +64,9 @@ impl<Provider: BlockHashReader> BlockHashReader for LatestStateProviderRef<'_, P
     }
 }
 
-impl<Provider: DBProvider> StateRootProvider for LatestStateProviderRef<'_, Provider> {
+impl<Provider: DBProvider + StateCommitmentProvider> StateRootProvider
+    for LatestStateProviderRef<'_, Provider>
+{
     fn state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
         StateRoot::overlay_root(self.tx(), hashed_state)
             .map_err(|err| ProviderError::Database(err.into()))
@@ -90,7 +94,9 @@ impl<Provider: DBProvider> StateRootProvider for LatestStateProviderRef<'_, Prov
     }
 }
 
-impl<Provider: DBProvider> StorageRootProvider for LatestStateProviderRef<'_, Provider> {
+impl<Provider: DBProvider + StateCommitmentProvider> StorageRootProvider
+    for LatestStateProviderRef<'_, Provider>
+{
     fn storage_root(
         &self,
         address: Address,
@@ -121,7 +127,9 @@ impl<Provider: DBProvider> StorageRootProvider for LatestStateProviderRef<'_, Pr
     }
 }
 
-impl<Provider: DBProvider> StateProofProvider for LatestStateProviderRef<'_, Provider> {
+impl<Provider: DBProvider + StateCommitmentProvider> StateProofProvider
+    for LatestStateProviderRef<'_, Provider>
+{
     fn proof(
         &self,
         input: TrieInput,
@@ -149,7 +157,7 @@ impl<Provider: DBProvider> StateProofProvider for LatestStateProviderRef<'_, Pro
     }
 }
 
-impl<Provider: DBProvider + BlockHashReader> StateProvider
+impl<Provider: DBProvider + BlockHashReader + StateCommitmentProvider> StateProvider
     for LatestStateProviderRef<'_, Provider>
 {
     /// Get storage.
@@ -191,7 +199,7 @@ impl<Provider: DBProvider> LatestStateProvider<Provider> {
 }
 
 // Delegates all provider impls to [LatestStateProviderRef]
-delegate_provider_impls!(LatestStateProvider<Provider> where [Provider: DBProvider + BlockHashReader]);
+delegate_provider_impls!(LatestStateProvider<Provider> where [Provider: DBProvider + BlockHashReader + StateCommitmentProvider]);
 
 #[cfg(test)]
 mod tests {
@@ -199,7 +207,9 @@ mod tests {
 
     const fn assert_state_provider<T: StateProvider>() {}
     #[allow(dead_code)]
-    const fn assert_latest_state_provider<T: DBProvider + BlockHashReader>() {
+    const fn assert_latest_state_provider<
+        T: DBProvider + BlockHashReader + StateCommitmentProvider,
+    >() {
         assert_state_provider::<LatestStateProvider<T>>();
     }
 }
diff --git a/crates/storage/storage-api/Cargo.toml b/crates/storage/storage-api/Cargo.toml
index c059eb0d6e9b..ba2ccf1b1573 100644
--- a/crates/storage/storage-api/Cargo.toml
+++ b/crates/storage/storage-api/Cargo.toml
@@ -23,6 +23,7 @@ reth-prune-types.workspace = true
 reth-stages-types.workspace = true
 reth-storage-errors.workspace = true
 reth-trie.workspace = true
+reth-trie-db.workspace = true
 reth-db.workspace = true
 
 # ethereum
diff --git a/crates/storage/storage-api/src/state.rs b/crates/storage/storage-api/src/state.rs
index 3174489fc4ac..0cb26d307434 100644
--- a/crates/storage/storage-api/src/state.rs
+++ b/crates/storage/storage-api/src/state.rs
@@ -8,6 +8,7 @@ use alloy_primitives::{Address, BlockHash, BlockNumber, StorageKey, StorageValue
 use auto_impl::auto_impl;
 use reth_primitives::Bytecode;
 use reth_storage_errors::provider::ProviderResult;
+use reth_trie_db::StateCommitment;
 
 /// Type alias of boxed [`StateProvider`].
 pub type StateProviderBox = Box<dyn StateProvider>;
@@ -81,6 +82,12 @@ pub trait StateProvider:
     }
 }
 
+/// Trait implemented for database providers that can provide the [`StateCommitment`] type.
+pub trait StateCommitmentProvider {
+    /// The [`StateCommitment`] type that can be used to perform state commitment operations.
+    type StateCommitment: StateCommitment;
+}
+
 /// Trait implemented for database providers that can be converted into a historical state provider.
 pub trait TryIntoHistoricalStateProvider {
     /// Returns a historical [`StateProvider`] indexed by the given historic block number.