From 59f75fbb4ce2785a4b5ed772e0107a20cd18f83b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BF=97=E5=AE=87?= <hello@evanlinjin.me>
Date: Fri, 15 Sep 2023 16:35:37 +0800
Subject: [PATCH] feat(wallet): add ability for wallet to record missing data

---
 crates/bdk/src/wallet/mod.rs                  | 65 ++++++++++++++++---
 .../wallet_esplora_async/src/main.rs          |  6 +-
 .../wallet_esplora_blocking/src/main.rs       |  6 +-
 3 files changed, 57 insertions(+), 20 deletions(-)

diff --git a/crates/bdk/src/wallet/mod.rs b/crates/bdk/src/wallet/mod.rs
index 94d566368c..4271cbe3fe 100644
--- a/crates/bdk/src/wallet/mod.rs
+++ b/crates/bdk/src/wallet/mod.rs
@@ -19,6 +19,7 @@ use alloc::{
     sync::Arc,
     vec::Vec,
 };
+use bdk_chain::collections::BTreeSet;
 pub use bdk_chain::keychain::Balance;
 use bdk_chain::{
     indexed_tx_graph,
@@ -91,6 +92,7 @@ pub struct Wallet<D = ()> {
     chain: LocalChain,
     indexed_graph: IndexedTxGraph<ConfirmationTimeAnchor, KeychainTxOutIndex<KeychainKind>>,
     persist: Persist<D, ChangeSet>,
+    missing: Missing,
     network: Network,
     secp: SecpCtx,
 }
@@ -209,6 +211,29 @@ impl From<keychain::ChangeSet<KeychainKind>> for ChangeSet {
     }
 }
 
+/// Represents data that is still missing after we [`apply_update`].
+///
+/// [`apply_update`]: Wallet::apply_update
+#[derive(Debug, Default, Clone)]
+#[must_use]
+pub struct Missing {
+    /// Heights of blocks that are missing in [`LocalChain`].
+    pub block_heights: BTreeSet<u32>,
+}
+
+impl Missing {
+    /// Whether there is missing data.
+    pub fn is_empty(&self) -> bool {
+        self.block_heights.is_empty()
+    }
+}
+
+impl fmt::Display for Missing {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "missing blocks={:?}", self.block_heights)
+    }
+}
+
 /// The address index selection strategy to use to derived an address from the wallet's external
 /// descriptor. See [`Wallet::get_address`]. If you're unsure which one to use use `WalletIndex::New`.
 #[derive(Debug)]
@@ -356,6 +381,14 @@ impl<D> Wallet<D> {
 
         let changeset = db.load_from_persistence().map_err(NewError::Persist)?;
         chain.apply_changeset(&changeset.chain);
+
+        let missing = Missing {
+            block_heights: changeset
+                .indexed_tx_graph
+                .graph
+                .missing_heights_from(&chain)
+                .collect(),
+        };
         indexed_graph.apply_changeset(changeset.indexed_tx_graph);
 
         let persist = Persist::new(db);
@@ -367,6 +400,7 @@ impl<D> Wallet<D> {
             chain,
             indexed_graph,
             persist,
+            missing,
             secp,
         })
     }
@@ -1965,26 +1999,37 @@ impl<D> Wallet<D> {
     where
         D: PersistBackend<ChangeSet>,
     {
-        let mut changeset = match update.chain {
-            Some(chain_update) => ChangeSet::from(self.chain.apply_update(chain_update)?),
-            None => ChangeSet::default(),
-        };
+        let mut changeset = ChangeSet::default();
+
+        if let Some(chain_update) = update.chain {
+            let chain_changeset = self.chain.apply_update(chain_update)?;
+            for height in chain_changeset.keys() {
+                self.missing.block_heights.remove(height);
+            }
+            changeset.append(ChangeSet::from(chain_changeset));
+        }
 
         let (_, index_changeset) = self
             .indexed_graph
             .index
             .reveal_to_target_multi(&update.last_active_indices);
-        changeset.append(ChangeSet::from(indexed_tx_graph::ChangeSet::from(
-            index_changeset,
-        )));
-        changeset.append(ChangeSet::from(
-            self.indexed_graph.apply_update(update.graph),
-        ));
+        changeset.append(ChangeSet::from(index_changeset));
+
+        let graph_changeset = self.indexed_graph.apply_update(update.graph);
+        self.missing
+            .block_heights
+            .extend(graph_changeset.graph.missing_heights_from(&self.chain));
+        changeset.append(ChangeSet::from(graph_changeset));
 
         self.persist.stage(changeset);
         Ok(())
     }
 
+    /// Get missing update data.
+    pub fn missing(&self) -> &Missing {
+        &self.missing
+    }
+
     /// Commits all currently [`staged`] changed to the persistence backend returning and error when
     /// this fails.
     ///
diff --git a/example-crates/wallet_esplora_async/src/main.rs b/example-crates/wallet_esplora_async/src/main.rs
index 5f57edc522..47758f104e 100644
--- a/example-crates/wallet_esplora_async/src/main.rs
+++ b/example-crates/wallet_esplora_async/src/main.rs
@@ -61,11 +61,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         graph: update_graph,
         ..Default::default()
     })?;
-    let missing_heights = wallet
-        .staged()
-        .indexed_tx_graph
-        .graph
-        .missing_heights_from(wallet.local_chain());
+    let missing_heights = wallet.missing().block_heights.iter().copied();
     let chain_update = client.update_local_chain(prev_tip, missing_heights).await?;
     wallet.apply_update(chain_update.into())?;
     wallet.commit()?;
diff --git a/example-crates/wallet_esplora_blocking/src/main.rs b/example-crates/wallet_esplora_blocking/src/main.rs
index 17c331360b..35f1d551e2 100644
--- a/example-crates/wallet_esplora_blocking/src/main.rs
+++ b/example-crates/wallet_esplora_blocking/src/main.rs
@@ -60,11 +60,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         last_active_indices,
         ..Default::default()
     })?;
-    let missing_heights = wallet
-        .staged()
-        .indexed_tx_graph
-        .graph
-        .missing_heights_from(wallet.local_chain());
+    let missing_heights = wallet.missing().block_heights.iter().copied();
     let chain_update = client.update_local_chain(prev_tip, missing_heights)?;
     wallet.apply_update(chain_update.into())?;
     wallet.commit()?;