Disconnect blocks from Wallet

crates/wallet/src/wallet/mod.rs

@@ -1081,6 +1081,23 @@ impl Wallet {
         Ok(changed)
     }
 
+    /// Disconnect a checkpoint and all checkpoints after it from the wallet's internal view of the chain.
+    ///
+    /// Returns whether anything changed after the disconnection. (e.g `false` if the block was
+    /// not present).
+    ///
+    /// **WARNING**: You must persist the changes resulting from one or more calls to this method
+    /// if you need the inserted changeset data to be reloaded after closing the wallet.
+    pub fn disconnect_checkpoint(&mut self, block_id: BlockId) -> bool {
+        if let Ok(changeset) = self.chain.disconnect_from(block_id) {
+            if !changeset.is_empty() {
+                self.stage.merge(changeset.into());
+                return true;
+            }
+        }
+        false
+    }
+
     /// Add a transaction to the wallet's internal view of the chain. This stages the change,
     /// you must persist it later.
     ///

crates/wallet/tests/wallet.rs

@@ -4308,3 +4308,25 @@ fn test_transactions_sort_by() {
         .collect();
     assert_eq!([None, Some(2000), Some(1000)], conf_heights.as_slice());
 }
+
+#[test]
+fn test_disconnect_applies() {
+    let (mut wallet, _txid) = get_funded_wallet_wpkh();
+    // We reorg a block
+    let checkpoint = wallet.local_chain().get(2_000).expect("existing block");
+    assert!(wallet.disconnect_checkpoint(checkpoint.block_id()));
+    let wallet_tip = wallet.local_chain().tip().block_id();
+    assert_ne!(wallet_tip.height, 2_000);
+    // Then we add some back
+    wallet
+        .insert_checkpoint(BlockId::from((2_500, BlockHash::all_zeros())))
+        .expect("valid block");
+    wallet
+        .insert_checkpoint(BlockId::from((3_000, BlockHash::all_zeros())))
+        .expect("valid block");
+    let wallet_tip = wallet.local_chain().tip().block_id();
+    let mut local_chain_iter = wallet.local_chain().iter_checkpoints();
+    let has_reorged = local_chain_iter.any(|checkpoint| checkpoint.height() == 2_000);
+    assert!(!has_reorged);
+    assert_eq!(wallet_tip.height, 3_000);
+}