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 "..".
This commit is contained in:
LLFourn 2024-02-06 17:31:22 +11:00 committed by 志宇
parent ee21ffeee0
commit fac228337c
No known key found for this signature in database
GPG Key ID: F6345C9837C2BDE8
6 changed files with 94 additions and 78 deletions

View File

@ -1015,7 +1015,7 @@ impl<D> Wallet<D> {
/// let (sent, received) = wallet.sent_and_received(tx); /// let (sent, received) = wallet.sent_and_received(tx);
/// ``` /// ```
pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) { 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). /// Get a single transaction from the wallet as a [`CanonicalTx`] (if the transaction exists).

View File

@ -268,15 +268,14 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
self.inner.unmark_used(&(keychain, index)) self.inner.unmark_used(&(keychain, index))
} }
/// Computes total input value going from script pubkeys in the index (sent) and the total output /// Computes the total value transfer effect `tx` has on the script pubkeys belonging to the
/// value going to script pubkeys in the index (received) in `tx`. For the `sent` to be computed /// keychains in `range`. Value is *sent* when a script pubkey in the `range` is on an input and
/// correctly, the output being spent must have already been scanned by the index. Calculating /// *received* when it is on an output. For `sent` to be computed correctly, the output being
/// received just uses the [`Transaction`] outputs directly, so it will be correct even if it has /// spent must have already been scanned by the index. Calculating received just uses the
/// not been scanned. /// [`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<K>) -> (u64, u64) {
/// This calls [`SpkTxOutIndex::sent_and_received`] internally. self.inner
pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) { .sent_and_received(tx, Self::map_to_inner_bounds(range))
self.inner.sent_and_received(tx)
} }
/// Computes the net value that this transaction gives to the script pubkeys in the index and /// Computes the net value that this transaction gives to the script pubkeys in the index and
@ -286,8 +285,8 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
/// This calls [`SpkTxOutIndex::net_value`] internally. /// This calls [`SpkTxOutIndex::net_value`] internally.
/// ///
/// [`sent_and_received`]: Self::sent_and_received /// [`sent_and_received`]: Self::sent_and_received
pub fn net_value(&self, tx: &Transaction) -> i64 { pub fn net_value(&self, tx: &Transaction, range: impl RangeBounds<K>) -> i64 {
self.inner.net_value(tx) self.inner.net_value(tx, Self::map_to_inner_bounds(range))
} }
} }
@ -390,24 +389,32 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
.collect() .collect()
} }
/// Iterate over revealed spks of all keychains. /// Iterate over revealed spks of keychains in `range`
pub fn revealed_spks(&self) -> impl DoubleEndedIterator<Item = (K, u32, &Script)> + Clone { pub fn revealed_spks(
self.keychains.keys().flat_map(|keychain| { &self,
self.revealed_keychain_spks(keychain) range: impl RangeBounds<K>,
.map(|(i, spk)| (keychain.clone(), i, spk)) ) -> impl DoubleEndedIterator<Item = (&K, u32, &Script)> + 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`. /// Iterate over revealed spks of the given `keychain`.
pub fn revealed_keychain_spks( pub fn revealed_keychain_spks<'a>(
&self, &'a self,
keychain: &K, keychain: &'a K,
) -> impl DoubleEndedIterator<Item = (u32, &Script)> + Clone { ) -> impl DoubleEndedIterator<Item = (u32, &Script)> + 'a {
let next_i = self.last_revealed.get(keychain).map_or(0, |&i| i + 1); self.revealed_spks(keychain..=keychain)
self.inner .map(|(_, i, spk)| (i, spk))
.all_spks()
.range((keychain.clone(), u32::MIN)..(keychain.clone(), next_i))
.map(|((_, i), spk)| (*i, spk.as_script()))
} }
/// Iterate over revealed, but unused, spks of all keychains. /// Iterate over revealed, but unused, spks of all keychains.
@ -617,38 +624,40 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
} }
} }
/// 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`. /// `keychain`.
/// pub fn keychain_outpoints<'a>(
/// Use [`keychain_outpoints_in_range`](KeychainTxOutIndex::keychain_outpoints_in_range) to &'a self,
/// iterate over a specific derivation range. keychain: &'a K,
pub fn keychain_outpoints( ) -> impl DoubleEndedIterator<Item = (u32, OutPoint)> + 'a {
&self, self.keychain_outpoints_in_range(keychain..=keychain)
keychain: &K, .map(move |(_, i, op)| (i, op))
) -> impl DoubleEndedIterator<Item = (u32, OutPoint)> + '_ {
self.keychain_outpoints_in_range(keychain, ..)
} }
/// Iterate over [`OutPoint`]s that point to `TxOut`s with script pubkeys derived from /// Iterate over [`OutPoint`]s that have script pubkeys derived from keychains in `range`.
/// `keychain` in a given derivation `range`. pub fn keychain_outpoints_in_range<'a>(
pub fn keychain_outpoints_in_range( &'a self,
&self, range: impl RangeBounds<K> + 'a,
keychain: &K, ) -> impl DoubleEndedIterator<Item = (&'a K, u32, OutPoint)> + 'a {
range: impl RangeBounds<u32>, let bounds = Self::map_to_inner_bounds(range);
) -> impl DoubleEndedIterator<Item = (u32, OutPoint)> + '_ {
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,
};
self.inner self.inner
.outputs_in_range((start, end)) .outputs_in_range(bounds)
.map(|((_, i), op)| (*i, op)) .map(move |((keychain, i), op)| (keychain, *i, op))
}
fn map_to_inner_bounds(bound: impl RangeBounds<K>) -> 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 /// Returns the highest derivation index of the `keychain` where [`KeychainTxOutIndex`] has

View File

@ -270,36 +270,39 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
self.spk_indices.get(script) self.spk_indices.get(script)
} }
/// Computes total input value going from script pubkeys in the index (sent) and the total output /// Computes the total value transfer effect `tx` has on the script pubkeys in `range`. Value is
/// value going to script pubkeys in the index (received) in `tx`. For the `sent` to be computed /// *sent* when a script pubkey in the `range` is on an input and *received* when it is on an
/// correctly, the output being spent must have already been scanned by the index. Calculating /// output. For `sent` to be computed correctly, the output being spent must have already been
/// received just uses the [`Transaction`] outputs directly, so it will be correct even if it has /// scanned by the index. Calculating received just uses the [`Transaction`] outputs directly,
/// not been scanned. /// so it will be correct even if it has not been scanned.
pub fn sent_and_received(&self, tx: &Transaction) -> (u64, u64) { pub fn sent_and_received(&self, tx: &Transaction, range: impl RangeBounds<I>) -> (u64, u64) {
let mut sent = 0; let mut sent = 0;
let mut received = 0; let mut received = 0;
for txin in &tx.input { for txin in &tx.input {
if let Some((_, txout)) = self.txout(txin.previous_output) { if let Some((index, txout)) = self.txout(txin.previous_output) {
sent += txout.value.to_sat(); if range.contains(index) {
sent += txout.value.to_sat();
}
} }
} }
for txout in &tx.output { for txout in &tx.output {
if self.index_of_spk(&txout.script_pubkey).is_some() { if let Some(index) = self.index_of_spk(&txout.script_pubkey) {
received += txout.value.to_sat(); if range.contains(index) {
received += txout.value.to_sat();
}
} }
} }
(sent, received) (sent, received)
} }
/// Computes the net value that this transaction gives to the script pubkeys in the index and /// Computes the net value transfer effect of `tx` on the script pubkeys in `range`. Shorthand
/// *takes* from the transaction outputs in the index. Shorthand for calling /// for calling [`sent_and_received`] and subtracting sent from received.
/// [`sent_and_received`] and subtracting sent from received.
/// ///
/// [`sent_and_received`]: Self::sent_and_received /// [`sent_and_received`]: Self::sent_and_received
pub fn net_value(&self, tx: &Transaction) -> i64 { pub fn net_value(&self, tx: &Transaction, range: impl RangeBounds<I>) -> i64 {
let (sent, received) = self.sent_and_received(tx); let (sent, received) = self.sent_and_received(tx, range);
received as i64 - sent as i64 received as i64 - sent as i64
} }

View File

@ -20,11 +20,13 @@ fn spk_txout_sent_and_received() {
}], }],
}; };
assert_eq!(index.sent_and_received(&tx1), (0, 42_000)); 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, ..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); index.index_tx(&tx1);
assert_eq!( assert_eq!(
index.sent_and_received(&tx1), index.sent_and_received(&tx1, ..),
(0, 42_000), (0, 42_000),
"shouldn't change after scanning" "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.sent_and_received(&tx2, ..), (42_000, 50_000));
assert_eq!(index.net_value(&tx2), 8_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] #[test]

View File

@ -210,8 +210,8 @@ fn main() -> anyhow::Result<()> {
if all_spks { if all_spks {
let all_spks = graph let all_spks = graph
.index .index
.revealed_spks() .revealed_spks(..)
.map(|(k, i, spk)| (k, i, spk.to_owned())) .map(|(k, i, spk)| (k.to_owned(), i, spk.to_owned()))
.collect::<Vec<_>>(); .collect::<Vec<_>>();
spks = Box::new(spks.chain(all_spks.into_iter().map(|(k, i, spk)| { spks = Box::new(spks.chain(all_spks.into_iter().map(|(k, i, spk)| {
eprintln!("scanning {}:{}", k, i); eprintln!("scanning {}:{}", k, i);

View File

@ -241,8 +241,8 @@ fn main() -> anyhow::Result<()> {
if *all_spks { if *all_spks {
let all_spks = graph let all_spks = graph
.index .index
.revealed_spks() .revealed_spks(..)
.map(|(k, i, spk)| (k, i, spk.to_owned())) .map(|(k, i, spk)| (k.to_owned(), i, spk.to_owned()))
.collect::<Vec<_>>(); .collect::<Vec<_>>();
spks = Box::new(spks.chain(all_spks.into_iter().map(|(k, i, spk)| { spks = Box::new(spks.chain(all_spks.into_iter().map(|(k, i, spk)| {
eprintln!("scanning {}:{}", k, i); eprintln!("scanning {}:{}", k, i);