From 022d481a7276fdd7a120452d45846bfaf2c0546e Mon Sep 17 00:00:00 2001
From: Dan Cline <6798349+Rjected@users.noreply.github.com>
Date: Tue, 9 Jul 2024 15:23:27 -0400
Subject: [PATCH] feat: add DatabaseProvider remove method (#9405)

---
 .../src/providers/database/provider.rs        | 26 ++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/crates/storage/provider/src/providers/database/provider.rs b/crates/storage/provider/src/providers/database/provider.rs
index 9e8e3bf355..e65be7682e 100644
--- a/crates/storage/provider/src/providers/database/provider.rs
+++ b/crates/storage/provider/src/providers/database/provider.rs
@@ -763,6 +763,22 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
         ))
     }
 
+    /// Remove list of entries from the table. Returns the number of entries removed.
+    #[inline]
+    pub fn remove<T: Table>(
+        &self,
+        range: impl RangeBounds<T::Key>,
+    ) -> Result<usize, DatabaseError> {
+        let mut entries = 0;
+        let mut cursor_write = self.tx.cursor_write::<T>()?;
+        let mut walker = cursor_write.walk_range(range)?;
+        while walker.next().transpose()?.is_some() {
+            walker.delete_current()?;
+            entries += 1;
+        }
+        Ok(entries)
+    }
+
     /// Return list of entries from table
     ///
     /// If TAKE is true, opened cursor would be write and it would delete all values from db.
@@ -894,7 +910,9 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
             // Remove TransactionBlocks index if there are transaction present
             if !transactions.is_empty() {
                 let tx_id_range = transactions.first().unwrap().0..=transactions.last().unwrap().0;
-                self.get_or_take::<tables::TransactionBlocks, TAKE>(tx_id_range)?;
+                // NOTE: we are in this branch because `TAKE` is true, so we can use the `remove`
+                // method
+                self.remove::<tables::TransactionBlocks>(tx_id_range)?;
             }
         }
 
@@ -959,7 +977,8 @@ impl<TX: DbTxMut + DbTx> DatabaseProvider<TX> {
 
         if TAKE {
             // rm HeaderTerminalDifficulties
-            self.get_or_take::<tables::HeaderTerminalDifficulties, TAKE>(range)?;
+            // NOTE: we are in this branch because `TAKE` is true, so we can use the `remove` method
+            self.remove::<tables::HeaderTerminalDifficulties>(range)?;
             // rm HeaderNumbers
             let mut header_number_cursor = self.tx.cursor_write::<tables::HeaderNumbers>()?;
             for (_, hash) in &block_header_hashes {
@@ -2612,7 +2631,8 @@ impl<TX: DbTxMut + DbTx> BlockExecutionWriter for DatabaseProvider<TX> {
         // that is why it is deleted afterwards.
         if TAKE {
             // rm block bodies
-            self.get_or_take::<tables::BlockBodyIndices, TAKE>(range)?;
+            // NOTE: we are in this branch because `TAKE` is true, so we can use the `remove` method
+            self.remove::<tables::BlockBodyIndices>(range)?;
 
             // Update pipeline progress
             if let Some(fork_number) = unwind_to {