From fac228337c79e1a7c406500c177e049e7f9f5da7 Mon Sep 17 00:00:00 2001 From: LLFourn Date: Tue, 6 Feb 2024 17:31:22 +1100 Subject: [PATCH] feat(chain)!: make `KeychainTxOutIndex` more range based `KeychainTxOutIndex` should try and avoid "all" kind of queries. There may be subranges of interest. If the user wants "all" they can just query "..". --- crates/bdk/src/wallet/mod.rs | 2 +- crates/chain/src/keychain/txout_index.rs | 115 +++++++++++--------- crates/chain/src/spk_txout_index.rs | 33 +++--- crates/chain/tests/test_spk_txout_index.rs | 14 ++- example-crates/example_electrum/src/main.rs | 4 +- example-crates/example_esplora/src/main.rs | 4 +- 6 files changed, 94 insertions(+), 78 deletions(-) diff --git a/crates/bdk/src/wallet/mod.rs b/crates/bdk/src/wallet/mod.rs index 84687882..f596f397 100644 --- a/crates/bdk/src/wallet/mod.rs +++ b/crates/bdk/src/wallet/mod.rs @@ -1015,7 +1015,7 @@ impl Wallet { /// let (sent, received) = wallet.sent_and_received(tx); /// ``` pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) { - self.indexed_graph.index.sent_and_received(tx) + self.indexed_graph.index.sent_and_received(tx, ..) } /// Get a single transaction from the wallet as a [`CanonicalTx`] (if the transaction exists). diff --git a/crates/chain/src/keychain/txout_index.rs b/crates/chain/src/keychain/txout_index.rs index 79f98fad..4b2479ec 100644 --- a/crates/chain/src/keychain/txout_index.rs +++ b/crates/chain/src/keychain/txout_index.rs @@ -268,15 +268,14 @@ impl KeychainTxOutIndex { self.inner.unmark_used(&(keychain, index)) } - /// Computes total input value going from script pubkeys in the index (sent) and the total output - /// value going to script pubkeys in the index (received) in `tx`. For the `sent` to be computed - /// correctly, the output being spent must have already been scanned by the index. Calculating - /// received just uses the [`Transaction`] outputs directly, so it will be correct even if it has - /// not been scanned. - /// - /// This calls [`SpkTxOutIndex::sent_and_received`] internally. - pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) { - self.inner.sent_and_received(tx) + /// Computes the total value transfer effect `tx` has on the script pubkeys belonging to the + /// keychains in `range`. Value is *sent* when a script pubkey in the `range` is on an input and + /// *received* when it is on an output. For `sent` to be computed correctly, the output being + /// spent must have already been scanned by the index. Calculating received just uses the + /// [`Transaction`] outputs directly, so it will be correct even if it has not been scanned. + pub fn sent_and_received(&self, tx: &Transaction, range: impl RangeBounds) -> (u64, u64) { + self.inner + .sent_and_received(tx, Self::map_to_inner_bounds(range)) } /// Computes the net value that this transaction gives to the script pubkeys in the index and @@ -286,8 +285,8 @@ impl KeychainTxOutIndex { /// This calls [`SpkTxOutIndex::net_value`] internally. /// /// [`sent_and_received`]: Self::sent_and_received - pub fn net_value(&self, tx: &Transaction) -> i64 { - self.inner.net_value(tx) + pub fn net_value(&self, tx: &Transaction, range: impl RangeBounds) -> i64 { + self.inner.net_value(tx, Self::map_to_inner_bounds(range)) } } @@ -390,24 +389,32 @@ impl KeychainTxOutIndex { .collect() } - /// Iterate over revealed spks of all keychains. - pub fn revealed_spks(&self) -> impl DoubleEndedIterator + Clone { - self.keychains.keys().flat_map(|keychain| { - self.revealed_keychain_spks(keychain) - .map(|(i, spk)| (keychain.clone(), i, spk)) + /// Iterate over revealed spks of keychains in `range` + pub fn revealed_spks( + &self, + range: impl RangeBounds, + ) -> impl DoubleEndedIterator + Clone { + self.keychains.range(range).flat_map(|(keychain, _)| { + let start = Bound::Included((keychain.clone(), u32::MIN)); + let end = match self.last_revealed.get(keychain) { + Some(last_revealed) => Bound::Included((keychain.clone(), *last_revealed)), + None => Bound::Excluded((keychain.clone(), u32::MIN)), + }; + + self.inner + .all_spks() + .range((start, end)) + .map(|((keychain, i), spk)| (keychain, *i, spk.as_script())) }) } /// Iterate over revealed spks of the given `keychain`. - pub fn revealed_keychain_spks( - &self, - keychain: &K, - ) -> impl DoubleEndedIterator + Clone { - let next_i = self.last_revealed.get(keychain).map_or(0, |&i| i + 1); - self.inner - .all_spks() - .range((keychain.clone(), u32::MIN)..(keychain.clone(), next_i)) - .map(|((_, i), spk)| (*i, spk.as_script())) + pub fn revealed_keychain_spks<'a>( + &'a self, + keychain: &'a K, + ) -> impl DoubleEndedIterator + 'a { + self.revealed_spks(keychain..=keychain) + .map(|(_, i, spk)| (i, spk)) } /// Iterate over revealed, but unused, spks of all keychains. @@ -617,38 +624,40 @@ impl KeychainTxOutIndex { } } - /// Iterate over all [`OutPoint`]s that point to `TxOut`s with script pubkeys derived from + /// Iterate over all [`OutPoint`]s that have `TxOut`s with script pubkeys derived from /// `keychain`. - /// - /// Use [`keychain_outpoints_in_range`](KeychainTxOutIndex::keychain_outpoints_in_range) to - /// iterate over a specific derivation range. - pub fn keychain_outpoints( - &self, - keychain: &K, - ) -> impl DoubleEndedIterator + '_ { - self.keychain_outpoints_in_range(keychain, ..) + pub fn keychain_outpoints<'a>( + &'a self, + keychain: &'a K, + ) -> impl DoubleEndedIterator + 'a { + self.keychain_outpoints_in_range(keychain..=keychain) + .map(move |(_, i, op)| (i, op)) } - /// Iterate over [`OutPoint`]s that point to `TxOut`s with script pubkeys derived from - /// `keychain` in a given derivation `range`. - pub fn keychain_outpoints_in_range( - &self, - keychain: &K, - range: impl RangeBounds, - ) -> impl DoubleEndedIterator + '_ { - let start = match range.start_bound() { - Bound::Included(i) => Bound::Included((keychain.clone(), *i)), - Bound::Excluded(i) => Bound::Excluded((keychain.clone(), *i)), - Bound::Unbounded => Bound::Unbounded, - }; - let end = match range.end_bound() { - Bound::Included(i) => Bound::Included((keychain.clone(), *i)), - Bound::Excluded(i) => Bound::Excluded((keychain.clone(), *i)), - Bound::Unbounded => Bound::Unbounded, - }; + /// Iterate over [`OutPoint`]s that have script pubkeys derived from keychains in `range`. + pub fn keychain_outpoints_in_range<'a>( + &'a self, + range: impl RangeBounds + 'a, + ) -> impl DoubleEndedIterator + 'a { + let bounds = Self::map_to_inner_bounds(range); self.inner - .outputs_in_range((start, end)) - .map(|((_, i), op)| (*i, op)) + .outputs_in_range(bounds) + .map(move |((keychain, i), op)| (keychain, *i, op)) + } + + fn map_to_inner_bounds(bound: impl RangeBounds) -> impl RangeBounds<(K, u32)> { + let start = match bound.start_bound() { + Bound::Included(keychain) => Bound::Included((keychain.clone(), u32::MIN)), + Bound::Excluded(keychain) => Bound::Excluded((keychain.clone(), u32::MAX)), + Bound::Unbounded => Bound::Unbounded, + }; + let end = match bound.end_bound() { + Bound::Included(keychain) => Bound::Included((keychain.clone(), u32::MAX)), + Bound::Excluded(keychain) => Bound::Excluded((keychain.clone(), u32::MIN)), + Bound::Unbounded => Bound::Unbounded, + }; + + (start, end) } /// Returns the highest derivation index of the `keychain` where [`KeychainTxOutIndex`] has diff --git a/crates/chain/src/spk_txout_index.rs b/crates/chain/src/spk_txout_index.rs index 24d06d12..a3ad48d0 100644 --- a/crates/chain/src/spk_txout_index.rs +++ b/crates/chain/src/spk_txout_index.rs @@ -270,36 +270,39 @@ impl SpkTxOutIndex { self.spk_indices.get(script) } - /// Computes total input value going from script pubkeys in the index (sent) and the total output - /// value going to script pubkeys in the index (received) in `tx`. For the `sent` to be computed - /// correctly, the output being spent must have already been scanned by the index. Calculating - /// received just uses the [`Transaction`] outputs directly, so it will be correct even if it has - /// not been scanned. - pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) { + /// Computes the total value transfer effect `tx` has on the script pubkeys in `range`. Value is + /// *sent* when a script pubkey in the `range` is on an input and *received* when it is on an + /// output. For `sent` to be computed correctly, the output being spent must have already been + /// scanned by the index. Calculating received just uses the [`Transaction`] outputs directly, + /// so it will be correct even if it has not been scanned. + pub fn sent_and_received(&self, tx: &Transaction, range: impl RangeBounds) -> (u64, u64) { let mut sent = 0; let mut received = 0; for txin in &tx.input { - if let Some((_, txout)) = self.txout(txin.previous_output) { - sent += txout.value.to_sat(); + if let Some((index, txout)) = self.txout(txin.previous_output) { + if range.contains(index) { + sent += txout.value.to_sat(); + } } } for txout in &tx.output { - if self.index_of_spk(&txout.script_pubkey).is_some() { - received += txout.value.to_sat(); + if let Some(index) = self.index_of_spk(&txout.script_pubkey) { + if range.contains(index) { + received += txout.value.to_sat(); + } } } (sent, received) } - /// Computes the net value that this transaction gives to the script pubkeys in the index and - /// *takes* from the transaction outputs in the index. Shorthand for calling - /// [`sent_and_received`] and subtracting sent from received. + /// Computes the net value transfer effect of `tx` on the script pubkeys in `range`. Shorthand + /// for calling [`sent_and_received`] and subtracting sent from received. /// /// [`sent_and_received`]: Self::sent_and_received - pub fn net_value(&self, tx: &Transaction) -> i64 { - let (sent, received) = self.sent_and_received(tx); + pub fn net_value(&self, tx: &Transaction, range: impl RangeBounds) -> i64 { + let (sent, received) = self.sent_and_received(tx, range); received as i64 - sent as i64 } diff --git a/crates/chain/tests/test_spk_txout_index.rs b/crates/chain/tests/test_spk_txout_index.rs index c84bd719..e6c39080 100644 --- a/crates/chain/tests/test_spk_txout_index.rs +++ b/crates/chain/tests/test_spk_txout_index.rs @@ -20,11 +20,13 @@ fn spk_txout_sent_and_received() { }], }; - assert_eq!(index.sent_and_received(&tx1), (0, 42_000)); - assert_eq!(index.net_value(&tx1), 42_000); + assert_eq!(index.sent_and_received(&tx1, ..), (0, 42_000)); + assert_eq!(index.sent_and_received(&tx1, ..1), (0, 42_000)); + assert_eq!(index.sent_and_received(&tx1, 1..), (0, 0)); + assert_eq!(index.net_value(&tx1, ..), 42_000); index.index_tx(&tx1); assert_eq!( - index.sent_and_received(&tx1), + index.sent_and_received(&tx1, ..), (0, 42_000), "shouldn't change after scanning" ); @@ -51,8 +53,10 @@ fn spk_txout_sent_and_received() { ], }; - assert_eq!(index.sent_and_received(&tx2), (42_000, 50_000)); - assert_eq!(index.net_value(&tx2), 8_000); + assert_eq!(index.sent_and_received(&tx2, ..), (42_000, 50_000)); + assert_eq!(index.sent_and_received(&tx2, ..1), (42_000, 30_000)); + assert_eq!(index.sent_and_received(&tx2, 1..), (0, 20_000)); + assert_eq!(index.net_value(&tx2, ..), 8_000); } #[test] diff --git a/example-crates/example_electrum/src/main.rs b/example-crates/example_electrum/src/main.rs index f651b85e..e3b758e7 100644 --- a/example-crates/example_electrum/src/main.rs +++ b/example-crates/example_electrum/src/main.rs @@ -210,8 +210,8 @@ fn main() -> anyhow::Result<()> { if all_spks { let all_spks = graph .index - .revealed_spks() - .map(|(k, i, spk)| (k, i, spk.to_owned())) + .revealed_spks(..) + .map(|(k, i, spk)| (k.to_owned(), i, spk.to_owned())) .collect::>(); spks = Box::new(spks.chain(all_spks.into_iter().map(|(k, i, spk)| { eprintln!("scanning {}:{}", k, i); diff --git a/example-crates/example_esplora/src/main.rs b/example-crates/example_esplora/src/main.rs index 9232081b..1abbf1ca 100644 --- a/example-crates/example_esplora/src/main.rs +++ b/example-crates/example_esplora/src/main.rs @@ -241,8 +241,8 @@ fn main() -> anyhow::Result<()> { if *all_spks { let all_spks = graph .index - .revealed_spks() - .map(|(k, i, spk)| (k, i, spk.to_owned())) + .revealed_spks(..) + .map(|(k, i, spk)| (k.to_owned(), i, spk.to_owned())) .collect::>(); spks = Box::new(spks.chain(all_spks.into_iter().map(|(k, i, spk)| { eprintln!("scanning {}:{}", k, i);