chore: move write_hashed_state from StorageWriter to trait StateChangeWriter (#9927)

crates/engine/tree/src/persistence.rs

@@ -6,8 +6,8 @@ use reth_errors::ProviderResult;
 use reth_primitives::{SealedBlock, StaticFileSegment, TransactionSignedNoHash, B256, U256};
 use reth_provider::{
     writer::StorageWriter, BlockExecutionWriter, BlockNumReader, BlockWriter, HistoryWriter,
-    OriginalValuesKnown, ProviderFactory, StageCheckpointWriter, StateWriter,
+    OriginalValuesKnown, ProviderFactory, StageCheckpointWriter, StateChangeWriter, StateWriter,
-    StaticFileProviderFactory, StaticFileWriter, TransactionsProviderExt,
+    StaticFileProviderFactory, StaticFileWriter, TransactionsProviderExt, TrieWriter,
 };
 use reth_prune::{Pruner, PrunerOutput};
 use reth_stages_types::{StageCheckpoint, StageId};
@@ -84,8 +84,8 @@ impl<DB: Database> PersistenceService<DB> {
             {
                 let trie_updates = block.trie_updates().clone();
                 let hashed_state = block.hashed_state();
-                storage_writer.write_hashed_state(&hashed_state.clone().into_sorted())?;
+                provider_rw.write_hashed_state(&hashed_state.clone().into_sorted())?;
-                storage_writer.write_trie_updates(&trie_updates)?;
+                provider_rw.write_trie_updates(&trie_updates)?;
             }
 
             // update history indices

crates/storage/provider/src/providers/database/provider.rs

@@ -2764,6 +2764,45 @@ impl<TX: DbTxMut + DbTx> StateChangeWriter for DatabaseProvider<TX> {
 
         Ok(())
     }
+
+    fn write_hashed_state(&self, hashed_state: &HashedPostStateSorted) -> ProviderResult<()> {
+        // Write hashed account updates.
+        let mut hashed_accounts_cursor = self.tx_ref().cursor_write::<tables::HashedAccounts>()?;
+        for (hashed_address, account) in hashed_state.accounts().accounts_sorted() {
+            if let Some(account) = account {
+                hashed_accounts_cursor.upsert(hashed_address, account)?;
+            } else if hashed_accounts_cursor.seek_exact(hashed_address)?.is_some() {
+                hashed_accounts_cursor.delete_current()?;
+            }
+        }
+
+        // Write hashed storage changes.
+        let sorted_storages = hashed_state.account_storages().iter().sorted_by_key(|(key, _)| *key);
+        let mut hashed_storage_cursor =
+            self.tx_ref().cursor_dup_write::<tables::HashedStorages>()?;
+        for (hashed_address, storage) in sorted_storages {
+            if storage.is_wiped() && hashed_storage_cursor.seek_exact(*hashed_address)?.is_some() {
+                hashed_storage_cursor.delete_current_duplicates()?;
+            }
+
+            for (hashed_slot, value) in storage.storage_slots_sorted() {
+                let entry = StorageEntry { key: hashed_slot, value };
+                if let Some(db_entry) =
+                    hashed_storage_cursor.seek_by_key_subkey(*hashed_address, entry.key)?
+                {
+                    if db_entry.key == entry.key {
+                        hashed_storage_cursor.delete_current()?;
+                    }
+                }
+
+                if !entry.value.is_zero() {
+                    hashed_storage_cursor.upsert(*hashed_address, entry)?;
+                }
+            }
+        }
+
+        Ok(())
+    }
 }
 
 impl<TX: DbTxMut + DbTx> TrieWriter for DatabaseProvider<TX> {
@@ -3532,6 +3571,7 @@ impl<DB: Database> BlockWriter for DatabaseProviderRW<DB> {
         Ok(block_indices)
     }
 
+    /// TODO(joshie): this fn should be moved to `StorageWriter` eventually
     fn append_blocks_with_state(
         &self,
         blocks: Vec<SealedBlockWithSenders>,
@@ -3566,10 +3606,8 @@ impl<DB: Database> BlockWriter for DatabaseProviderRW<DB> {
         durations_recorder.record_relative(metrics::Action::InsertState);
 
         // insert hashes and intermediate merkle nodes
-        {
+        self.write_hashed_state(&hashed_state)?;
-            storage_writer.write_hashed_state(&hashed_state)?;
+        self.write_trie_updates(&trie_updates)?;
-            self.write_trie_updates(&trie_updates)?;
-        }
         durations_recorder.record_relative(metrics::Action::InsertHashes);
 
         self.update_history_indices(first_number..=last_block_number)?;

crates/storage/provider/src/traits/state.rs

@@ -1,6 +1,7 @@
 use reth_execution_types::ExecutionOutcome;
 use reth_primitives::BlockNumber;
 use reth_storage_errors::provider::ProviderResult;
+use reth_trie::HashedPostStateSorted;
 use revm::db::{
     states::{PlainStateReverts, StateChangeset},
     OriginalValuesKnown,
@@ -30,4 +31,7 @@ pub trait StateChangeWriter {
 
     /// Write state changes to the database.
     fn write_state_changes(&self, changes: StateChangeset) -> ProviderResult<()>;
+
+    /// Writes the hashed state changes to the database
+    fn write_hashed_state(&self, hashed_state: &HashedPostStateSorted) -> ProviderResult<()>;
 }

crates/storage/provider/src/writer/mod.rs

@@ -2,9 +2,8 @@ use crate::{
     providers::StaticFileProviderRWRefMut, DatabaseProvider, DatabaseProviderRO,
     DatabaseProviderRW, StateChangeWriter, StateWriter, TrieWriter,
 };
-use itertools::Itertools;
 use reth_db::{
-    cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO, DbDupCursorRW},
+    cursor::DbCursorRO,
     tables,
     transaction::{DbTx, DbTxMut},
     Database,
@@ -12,11 +11,11 @@ use reth_db::{
 use reth_errors::{ProviderError, ProviderResult};
 use reth_execution_types::ExecutionOutcome;
 use reth_primitives::{
-    BlockNumber, Header, StaticFileSegment, StorageEntry, TransactionSignedNoHash, B256, U256,
+    BlockNumber, Header, StaticFileSegment, TransactionSignedNoHash, B256, U256,
 };
 use reth_storage_api::ReceiptWriter;
 use reth_storage_errors::writer::StorageWriterError;
-use reth_trie::{updates::TrieUpdates, HashedPostStateSorted};
+use reth_trie::updates::TrieUpdates;
 use revm::db::OriginalValuesKnown;
 use static_file::StaticFileWriter;
 use std::borrow::Borrow;
@@ -219,49 +218,6 @@ impl<'a, 'b, TX> StorageWriter<'a, 'b, TX>
 where
     TX: DbTxMut + DbTx,
 {
-    /// Writes the hashed state changes to the database
-    pub fn write_hashed_state(&self, hashed_state: &HashedPostStateSorted) -> ProviderResult<()> {
-        self.ensure_database_writer()?;
-
-        // Write hashed account updates.
-        let mut hashed_accounts_cursor =
-            self.database_writer().tx_ref().cursor_write::<tables::HashedAccounts>()?;
-        for (hashed_address, account) in hashed_state.accounts().accounts_sorted() {
-            if let Some(account) = account {
-                hashed_accounts_cursor.upsert(hashed_address, account)?;
-            } else if hashed_accounts_cursor.seek_exact(hashed_address)?.is_some() {
-                hashed_accounts_cursor.delete_current()?;
-            }
-        }
-
-        // Write hashed storage changes.
-        let sorted_storages = hashed_state.account_storages().iter().sorted_by_key(|(key, _)| *key);
-        let mut hashed_storage_cursor =
-            self.database_writer().tx_ref().cursor_dup_write::<tables::HashedStorages>()?;
-        for (hashed_address, storage) in sorted_storages {
-            if storage.is_wiped() && hashed_storage_cursor.seek_exact(*hashed_address)?.is_some() {
-                hashed_storage_cursor.delete_current_duplicates()?;
-            }
-
-            for (hashed_slot, value) in storage.storage_slots_sorted() {
-                let entry = StorageEntry { key: hashed_slot, value };
-                if let Some(db_entry) =
-                    hashed_storage_cursor.seek_by_key_subkey(*hashed_address, entry.key)?
-                {
-                    if db_entry.key == entry.key {
-                        hashed_storage_cursor.delete_current()?;
-                    }
-                }
-
-                if !entry.value.is_zero() {
-                    hashed_storage_cursor.upsert(*hashed_address, entry)?;
-                }
-            }
-        }
-
-        Ok(())
-    }
-
     /// Appends receipts block by block.
     ///
     /// ATTENTION: If called from [`StorageWriter`] without a static file producer, it will always
@@ -381,7 +337,7 @@ mod tests {
     use crate::{test_utils::create_test_provider_factory, AccountReader, TrieWriter};
     use reth_db::tables;
     use reth_db_api::{
-        cursor::{DbCursorRO, DbDupCursorRO},
+        cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO},
         models::{AccountBeforeTx, BlockNumberAddress},
         transaction::{DbTx, DbTxMut},
     };
@@ -437,8 +393,7 @@ mod tests {
         hashed_state.storages.insert(destroyed_address_hashed, HashedStorage::new(true));
 
         let provider_rw = provider_factory.provider_rw().unwrap();
-        let storage_writer = StorageWriter::new(Some(&provider_rw), None);
+        assert_eq!(provider_rw.write_hashed_state(&hashed_state.into_sorted()), Ok(()));
-        assert_eq!(storage_writer.write_hashed_state(&hashed_state.into_sorted()), Ok(()));
         provider_rw.commit().unwrap();
 
         let provider = provider_factory.provider().unwrap();

crates/trie/parallel/benches/root.rs

@@ -5,7 +5,7 @@ use proptest_arbitrary_interop::arb;
 use rayon::ThreadPoolBuilder;
 use reth_primitives::{Account, B256, U256};
 use reth_provider::{
-    providers::ConsistentDbView, test_utils::create_test_provider_factory, writer::StorageWriter,
+    providers::ConsistentDbView, test_utils::create_test_provider_factory, StateChangeWriter,
     TrieWriter,
 };
 use reth_tasks::pool::BlockingTaskPool;
@@ -29,8 +29,7 @@ pub fn calculate_state_root(c: &mut Criterion) {
         let provider_factory = create_test_provider_factory();
         {
             let provider_rw = provider_factory.provider_rw().unwrap();
-            let storage_writer = StorageWriter::new(Some(&provider_rw), None);
+            provider_rw.write_hashed_state(&db_state.into_sorted()).unwrap();
-            storage_writer.write_hashed_state(&db_state.into_sorted()).unwrap();
             let (_, updates) =
                 StateRoot::from_tx(provider_rw.tx_ref()).root_with_updates().unwrap();
             provider_rw.write_trie_updates(&updates).unwrap();