return rollback back

crates/exex/exex/src/wal/mod.rs

@@ -102,6 +102,79 @@ impl Wal {
         Ok(())
     }
 
+    /// Rollbacks the WAL to the given block, inclusive.
+    ///
+    /// 1. Walks the WAL from the end and searches for the first notification where committed chain
+    ///    contains a block with the same number and hash as `to_block`.
+    /// 2. If the notification is found, truncates the WAL. It means that if the found notification
+    ///    contains both given block and blocks before it, the whole notification will be truncated.
+    ///
+    /// # Returns
+    ///
+    /// 1. The block number and hash of the lowest removed block.
+    /// 2. The notifications that were removed.
+    #[instrument(target = "exex::wal", skip(self))]
+    pub fn rollback(
+        &mut self,
+        to_block: BlockNumHash,
+    ) -> eyre::Result<Option<(BlockNumHash, Vec<ExExNotification>)>> {
+        // First, pop items from the back of the cache until we find the notification with the
+        // specified block. When found, save the file ID of that notification.
+        let mut remove_from_file_id = None;
+        let mut remove_to_file_id = None;
+        let mut lowest_removed_block = None;
+        while let Some((file_id, block)) = self.block_cache.pop_back() {
+            debug!(?file_id, ?block, "Popped back block from the block cache");
+            if block.action.is_commit() && block.block.number == to_block.number {
+                debug!(
+                    ?file_id,
+                    ?block,
+                    ?remove_from_file_id,
+                    ?lowest_removed_block,
+                    "Found the requested block"
+                );
+
+                if block.block.hash != to_block.hash {
+                    eyre::bail!("block hash mismatch in WAL")
+                }
+
+                remove_from_file_id = Some(file_id);
+
+                let notification = self.storage.read_notification(file_id)?;
+                lowest_removed_block = notification
+                    .committed_chain()
+                    .as_ref()
+                    .map(|chain| chain.first())
+                    .map(|block| (block.number, block.hash()).into());
+
+                break
+            }
+
+            remove_from_file_id = Some(file_id);
+            remove_to_file_id.get_or_insert(file_id);
+        }
+
+        // If the specified block is still not found, we can't do anything and just return. The
+        // cache was empty.
+        let Some((remove_from_file_id, remove_to_file_id)) =
+            remove_from_file_id.zip(remove_to_file_id)
+        else {
+            debug!("No blocks were rolled back");
+            return Ok(None)
+        };
+
+        // Remove the rest of the block cache entries for the file ID that we found.
+        self.block_cache.remove_notification(remove_from_file_id);
+        debug!(?remove_from_file_id, "Block cache was rolled back");
+
+        // Remove notifications from the storage.
+        let removed_notifications =
+            self.storage.take_notifications(remove_from_file_id..=remove_to_file_id)?;
+        debug!(removed_notifications = ?removed_notifications.len(), "Storage was rolled back");
+
+        Ok(Some((lowest_removed_block.expect("qed"), removed_notifications)))
+    }
+
     /// Finalizes the WAL to the given block, inclusive.
     ///
     /// 1. Finds a notification with first unfinalized block (first notification containing a