Skip to content

Commit

Permalink
feat(wallet): add ability for wallet to record missing data
Browse files Browse the repository at this point in the history
  • Loading branch information
evanlinjin committed Sep 15, 2023
1 parent 8e82bfa commit 59f75fb
Show file tree
Hide file tree
Showing 3 changed files with 57 additions and 20 deletions.
65 changes: 55 additions & 10 deletions crates/bdk/src/wallet/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
}
Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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);
Expand All @@ -367,6 +400,7 @@ impl<D> Wallet<D> {
chain,
indexed_graph,
persist,
missing,
secp,
})
}
Expand Down Expand Up @@ -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.
///
Expand Down
6 changes: 1 addition & 5 deletions example-crates/wallet_esplora_async/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()?;
Expand Down
6 changes: 1 addition & 5 deletions example-crates/wallet_esplora_blocking/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()?;
Expand Down

0 comments on commit 59f75fb

Please sign in to comment.