From 119852fc8532dcbec6c45ff84e6ebbd4aff076ff Mon Sep 17 00:00:00 2001
From: Santiago Carmuega <santiago@carmuega.me>
Date: Sun, 10 Sep 2023 14:45:24 +0200
Subject: [PATCH] refactor(rolldb): use iterator for chain rollback delete
 (#85)

---
 src/storage/rolldb/mod.rs | 32 ++++++++++++++++++++++++++------
 src/storage/rolldb/wal.rs |  6 ++----
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/storage/rolldb/mod.rs b/src/storage/rolldb/mod.rs
index 70b6efc1..c5228a84 100644
--- a/src/storage/rolldb/mod.rs
+++ b/src/storage/rolldb/mod.rs
@@ -91,12 +91,13 @@ impl RollDB {
     pub fn roll_back(&mut self, until: BlockSlot) -> Result<(), Error> {
         let mut batch = WriteBatch::default();
 
-        let (new_seq, removed_blocks) =
-            wal::WalKV::stage_roll_back(&self.db, self.wal_seq, until, &mut batch)?;
+        let new_seq = wal::WalKV::stage_roll_back(&self.db, self.wal_seq, until, &mut batch)?;
 
-        // remove rollbacked blocks from ChainKV
-        for slot in removed_blocks {
-            ChainKV::stage_delete(&self.db, DBInt(slot), &mut batch);
+        // remove rollback-ed blocks from ChainKV
+        let to_remove = ChainKV::iter_keys_from(&self.db, DBInt(until)).skip(1);
+
+        for key in to_remove {
+            ChainKV::stage_delete(&self.db, key?, &mut batch);
         }
 
         self.db.write(batch).map_err(|_| Error::IO)?;
@@ -371,27 +372,46 @@ mod tests {
             let (slot, hash, body) = dummy_block(11);
             db.roll_forward(slot, hash, body.clone()).unwrap();
 
+            // ensure block body is persisted
             let persisted = db.get_block(hash).unwrap().unwrap();
             assert_eq!(persisted, body);
 
+            // ensure tip matches
             let (tip_slot, tip_hash) = db.find_tip().unwrap().unwrap();
             assert_eq!(tip_slot, slot);
             assert_eq!(tip_hash, hash);
+
+            // ensure chain has item
+            let (chain_slot, chain_hash) = db.crawl_chain().next().unwrap().unwrap();
+            assert_eq!(chain_slot, slot);
+            assert_eq!(chain_hash, hash);
         });
     }
 
     #[test]
     fn test_roll_back_blackbox() {
         with_tmp_db(30, |mut db| {
-            for i in 0..5 {
+            for i in 0..=5 {
                 let (slot, hash, body) = dummy_block(i * 10);
                 db.roll_forward(slot, hash, body).unwrap();
             }
 
             db.roll_back(20).unwrap();
 
+            // ensure tip show rollback point
             let (tip_slot, _) = db.find_tip().unwrap().unwrap();
             assert_eq!(tip_slot, 20);
+
+            // ensure chain has items not rolled back
+            let mut chain = db.crawl_chain();
+
+            for i in 0..=2 {
+                let (slot, _) = chain.next().unwrap().unwrap();
+                assert_eq!(slot, i * 10);
+            }
+
+            // ensure chain stops here
+            assert!(chain.next().is_none());
         });
     }
 
diff --git a/src/storage/rolldb/wal.rs b/src/storage/rolldb/wal.rs
index 1e650cfa..ea860c08 100644
--- a/src/storage/rolldb/wal.rs
+++ b/src/storage/rolldb/wal.rs
@@ -90,9 +90,8 @@ impl WalKV {
         mut last_seq: Seq,
         until: super::BlockSlot,
         batch: &mut WriteBatch,
-    ) -> Result<(u64, Vec<BlockSlot>), super::Error> {
+    ) -> Result<u64, super::Error> {
         let iter = WalKV::iter_values(db, IteratorMode::End);
-        let mut removed_blocks = vec![];
 
         for step in iter {
             let value = step.map_err(|_| super::Error::IO)?.0;
@@ -104,14 +103,13 @@ impl WalKV {
 
             match value.into_undo() {
                 Some(undo) => {
-                    removed_blocks.push(undo.slot());
                     last_seq = Self::stage_append(db, last_seq, undo, batch)?;
                 }
                 None => continue,
             };
         }
 
-        Ok((last_seq, removed_blocks))
+        Ok(last_seq)
     }
 
     pub fn stage_roll_forward(