feat(prune): Add MerkleChangeSets segment

Author: RomanHodulak

crates/prune/prune/src/segments/mod.rs

@@ -15,8 +15,8 @@ pub use static_file::{
 use std::{fmt::Debug, ops::RangeInclusive};
 use tracing::error;
 pub use user::{
-    AccountHistory, Receipts as UserReceipts, ReceiptsByLogs, SenderRecovery, StorageHistory,
-    TransactionLookup,
+    AccountHistory, MerkleChangeSets, Receipts as UserReceipts, ReceiptsByLogs, SenderRecovery,
+    StorageHistory, TransactionLookup,
 };
 
 /// A segment represents a pruning of some portion of the data.

crates/prune/prune/src/segments/user/merkle_change_sets.rs

@@ -0,0 +1,101 @@
+use crate::{
+    db_ext::DbTxPruneExt,
+    segments::{PruneInput, Segment},
+    PrunerError,
+};
+use alloy_primitives::B256;
+use reth_db_api::{models::BlockNumberHashedAddress, table::Value, tables, transaction::DbTxMut};
+use reth_primitives_traits::NodePrimitives;
+use reth_provider::{
+    errors::provider::ProviderResult, BlockReader, DBProvider, NodePrimitivesProvider,
+    PruneCheckpointWriter, TransactionsProvider,
+};
+use reth_prune_types::{
+    PruneCheckpoint, PruneMode, PrunePurpose, PruneSegment, SegmentOutput, SegmentOutputCheckpoint,
+};
+use tracing::{instrument, trace};
+
+#[derive(Debug)]
+pub struct MerkleChangeSets {
+    mode: PruneMode,
+}
+
+impl MerkleChangeSets {
+    pub const fn new(mode: PruneMode) -> Self {
+        Self { mode }
+    }
+}
+
+impl<Provider> Segment<Provider> for MerkleChangeSets
+where
+    Provider: DBProvider<Tx: DbTxMut>
+        + PruneCheckpointWriter
+        + TransactionsProvider
+        + BlockReader
+        + NodePrimitivesProvider<Primitives: NodePrimitives<Receipt: Value>>,
+{
+    fn segment(&self) -> PruneSegment {
+        PruneSegment::MerkleChangeSets
+    }
+
+    fn mode(&self) -> Option<PruneMode> {
+        Some(self.mode)
+    }
+
+    fn purpose(&self) -> PrunePurpose {
+        PrunePurpose::User
+    }
+
+    #[instrument(level = "trace", target = "pruner", skip(self, provider), ret)]
+    fn prune(&self, provider: &Provider, input: PruneInput) -> Result<SegmentOutput, PrunerError> {
+        let Some(block_range) = input.get_next_block_range() else {
+            trace!(target: "pruner", "No change sets to prune");
+            return Ok(SegmentOutput::done())
+        };
+
+        let block_range_end = *block_range.end();
+        let mut limiter = input.limiter;
+        let mut last_pruned_block = block_range_end;
+
+        for block in block_range {
+            let (pruned, done) =
+                provider.tx_ref().prune_table_with_range::<tables::StoragesTrieChangeSets>(
+                    BlockNumberHashedAddress::from((block, B256::ZERO))..,
+                    &mut limiter,
+                    |_| false,
+                    |row| last_pruned_block = row.0 .0 .0,
+                )?;
+
+            trace!(target: "pruner", %pruned, %done, "Pruned storages change sets");
+        }
+
+        let (pruned, done) =
+            provider.tx_ref().prune_table_with_range::<tables::AccountsTrieChangeSets>(
+                block_range.clone(),
+                &mut limiter,
+                |_| false,
+                |row| last_pruned_block = row.0,
+            )?;
+
+        trace!(target: "pruner", %pruned, %done, "Pruned accounts change sets");
+
+        let progress = limiter.progress(done);
+
+        Ok(SegmentOutput {
+            progress,
+            pruned,
+            checkpoint: Some(SegmentOutputCheckpoint {
+                block_number: Some(last_pruned_block),
+                tx_number: None,
+            }),
+        })
+    }
+
+    fn save_checkpoint(
+        &self,
+        provider: &Provider,
+        checkpoint: PruneCheckpoint,
+    ) -> ProviderResult<()> {
+        provider.save_prune_checkpoint(PruneSegment::MerkleChangeSets, checkpoint)
+    }
+}

crates/prune/prune/src/segments/user/mod.rs

@@ -1,12 +1,14 @@
 mod account_history;
 mod history;
+mod merkle_change_sets;
 mod receipts;
 mod receipts_by_logs;
 mod sender_recovery;
 mod storage_history;
 mod transaction_lookup;
 
 pub use account_history::AccountHistory;
+pub use merkle_change_sets::MerkleChangeSets;
 pub use receipts::Receipts;
 pub use receipts_by_logs::ReceiptsByLogs;
 pub use sender_recovery::SenderRecovery;

crates/prune/types/src/segment.rs

@@ -15,6 +15,9 @@ pub enum PruneSegment {
     TransactionLookup,
     /// Prune segment responsible for all rows in `Receipts` table.
     Receipts,
+    /// Prune segment responsible for all rows in `AccountsTrieChangeSets` and
+    /// `StoragesTrieChangeSets` table.
+    MerkleChangeSets,
     /// Prune segment responsible for some rows in `Receipts` table filtered by logs.
     ContractLogs,
     /// Prune segment responsible for the `AccountChangeSets` and `AccountsHistory` tables.
@@ -44,9 +47,10 @@ impl PruneSegment {
                 0
             }
             Self::Receipts if purpose.is_static_file() => 0,
-            Self::ContractLogs | Self::AccountHistory | Self::StorageHistory => {
-                MINIMUM_PRUNING_DISTANCE
-            }
+            Self::ContractLogs |
+            Self::AccountHistory |
+            Self::StorageHistory |
+            Self::MerkleChangeSets |
             Self::Receipts => MINIMUM_PRUNING_DISTANCE,
         }
     }

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

@@ -2272,7 +2272,7 @@ mod tests {
 
             // Invalid/Non-existent argument should return `None`
             {
-                call_method!($arg_count, provider, $method, |_,_,_,_| ( ($invalid_args, None)), tx_num, tx_hash, &in_memory_blocks[0], &receipts);
+                call_method!($arg_count, provider, $method, |_,_,_,_| ($invalid_args, None), tx_num, tx_hash, &in_memory_blocks[0], &receipts);
             }
 
             // Check that the item is only in memory and not in database
@@ -2283,7 +2283,7 @@ mod tests {
                 call_method!($arg_count, provider, $method, |_,_,_,_| (args.clone(), expected_item), tx_num, tx_hash, last_mem_block, &receipts);
 
                 // Ensure the item is not in storage
-                call_method!($arg_count, provider.database, $method, |_,_,_,_| ( (args, None)), tx_num, tx_hash, last_mem_block, &receipts);
+                call_method!($arg_count, provider.database, $method, |_,_,_,_| (args, None), tx_num, tx_hash, last_mem_block, &receipts);
             }
         )*
     }};