From 0509f11644849fb3279536d380654a6b31cf9a7c 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                  | 52 +++++++++++++++----
 .../wallet_esplora_async/src/main.rs          |  6 +--
 .../wallet_esplora_blocking/src/main.rs       |  6 +--
 3 files changed, 44 insertions(+), 20 deletions(-)

diff --git a/crates/bdk/src/wallet/mod.rs b/crates/bdk/src/wallet/mod.rs
index 94d566368..d27653a0c 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,10 +92,22 @@ pub struct Wallet<D = ()> {
     chain: LocalChain,
     indexed_graph: IndexedTxGraph<ConfirmationTimeAnchor, KeychainTxOutIndex<KeychainKind>>,
     persist: Persist<D, ChangeSet>,
+    missing_heights: MissingBlockHeights,
     network: Network,
     secp: SecpCtx,
 }
 
+/// Heights of blocks that ae missing from the wallet's [`LocalChain`].
+///
+/// Some chain sources may fetch transactions that are [`Anchor`]ed to block heights that do not
+/// exist in our [`LocalChain`]. A second call is needed to fetch those missing [`BlockId`]
+/// checkpoints.
+///
+/// Refer to [`Wallet::missing_heights`] for more.
+///
+/// [`Anchor`]: bdk_chain::Anchor
+pub type MissingBlockHeights = BTreeSet<u32>;
+
 /// An update to [`Wallet`].
 ///
 /// It updates [`bdk_chain::keychain::KeychainTxOutIndex`], [`bdk_chain::TxGraph`] and [`local_chain::LocalChain`] atomically.
@@ -356,6 +369,12 @@ impl<D> Wallet<D> {
 
         let changeset = db.load_from_persistence().map_err(NewError::Persist)?;
         chain.apply_changeset(&changeset.chain);
+
+        let missing_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 +386,7 @@ impl<D> Wallet<D> {
             chain,
             indexed_graph,
             persist,
+            missing_heights,
             secp,
         })
     }
@@ -1965,26 +1985,38 @@ 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_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_heights
+            .extend(graph_changeset.graph.missing_heights_from(&self.chain));
+        changeset.append(ChangeSet::from(graph_changeset));
 
         self.persist.stage(changeset);
         Ok(())
     }
 
+    /// Data that is still missing after we call [`Wallet::apply_update`].
+    ///
+    /// Some chain sources requires multiple rounds of I/O.
+    pub fn missing_heights(&self) -> &MissingBlockHeights {
+        &self.missing_heights
+    }
+
     /// 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 5f57edc52..d4a435203 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_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 17c331360..264ef15df 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_heights().iter().copied();
     let chain_update = client.update_local_chain(prev_tip, missing_heights)?;
     wallet.apply_update(chain_update.into())?;
     wallet.commit()?;