Improve txout listing and balance APIs for redesigned structures

Instead of relying on a `OwnedIndexer` trait to filter for relevant
txouts, we move the listing and balance methods from `IndexedTxGraph` to
`TxGraph` and add an additional input (list of relevant outpoints) to
these methods.

This provides a simpler implementation and a more flexible API.

Author: evanlinjin

crates/chain/src/indexed_tx_graph.rs

@@ -1,12 +1,9 @@
-use core::convert::Infallible;
-
 use alloc::vec::Vec;
-use bitcoin::{OutPoint, Script, Transaction, TxOut};
+use bitcoin::{OutPoint, Transaction, TxOut};
 
 use crate::{
-    keychain::Balance,
     tx_graph::{Additions, TxGraph},
-    Anchor, Append, BlockId, ChainOracle, FullTxOut, ObservedAs,
+    Anchor, Append,
 };
 
 /// A struct that combines [`TxGraph`] and an [`Indexer`] implementation.
@@ -29,6 +26,14 @@ impl<A, I: Default> Default for IndexedTxGraph<A, I> {
 }
 
 impl<A, I> IndexedTxGraph<A, I> {
+    /// Construct a new [`IndexedTxGraph`] with a given `index`.
+    pub fn new(index: I) -> Self {
+        Self {
+            index,
+            graph: TxGraph::default(),
+        }
+    }
+
     /// Get a reference of the internal transaction graph.
     pub fn graph(&self) -> &TxGraph<A> {
         &self.graph
@@ -157,115 +162,6 @@ where
     }
 }
 
-impl<A: Anchor, I: OwnedIndexer> IndexedTxGraph<A, I> {
-    pub fn try_list_owned_txouts<'a, C: ChainOracle + 'a>(
-        &'a self,
-        chain: &'a C,
-        chain_tip: BlockId,
-    ) -> impl Iterator<Item = Result<FullTxOut<ObservedAs<A>>, C::Error>> + 'a {
-        self.graph()
-            .try_list_chain_txouts(chain, chain_tip)
-            .filter(|r| {
-                if let Ok(full_txout) = r {
-                    if !self.index.is_spk_owned(&full_txout.txout.script_pubkey) {
-                        return false;
-                    }
-                }
-                true
-            })
-    }
-
-    pub fn list_owned_txouts<'a, C: ChainOracle<Error = Infallible> + 'a>(
-        &'a self,
-        chain: &'a C,
-        chain_tip: BlockId,
-    ) -> impl Iterator<Item = FullTxOut<ObservedAs<A>>> + 'a {
-        self.try_list_owned_txouts(chain, chain_tip)
-            .map(|r| r.expect("oracle is infallible"))
-    }
-
-    pub fn try_list_owned_unspents<'a, C: ChainOracle + 'a>(
-        &'a self,
-        chain: &'a C,
-        chain_tip: BlockId,
-    ) -> impl Iterator<Item = Result<FullTxOut<ObservedAs<A>>, C::Error>> + 'a {
-        self.graph()
-            .try_list_chain_unspents(chain, chain_tip)
-            .filter(|r| {
-                if let Ok(full_txout) = r {
-                    if !self.index.is_spk_owned(&full_txout.txout.script_pubkey) {
-                        return false;
-                    }
-                }
-                true
-            })
-    }
-
-    pub fn list_owned_unspents<'a, C: ChainOracle<Error = Infallible> + 'a>(
-        &'a self,
-        chain: &'a C,
-        chain_tip: BlockId,
-    ) -> impl Iterator<Item = FullTxOut<ObservedAs<A>>> + 'a {
-        self.try_list_owned_unspents(chain, chain_tip)
-            .map(|r| r.expect("oracle is infallible"))
-    }
-
-    pub fn try_balance<C, F>(
-        &self,
-        chain: &C,
-        chain_tip: BlockId,
-        mut should_trust: F,
-    ) -> Result<Balance, C::Error>
-    where
-        C: ChainOracle,
-        F: FnMut(&Script) -> bool,
-    {
-        let tip_height = chain_tip.height;
-
-        let mut immature = 0;
-        let mut trusted_pending = 0;
-        let mut untrusted_pending = 0;
-        let mut confirmed = 0;
-
-        for res in self.try_list_owned_unspents(chain, chain_tip) {
-            let txout = res?;
-
-            match &txout.chain_position {
-                ObservedAs::Confirmed(_) => {
-                    if txout.is_confirmed_and_spendable(tip_height) {
-                        confirmed += txout.txout.value;
-                    } else if !txout.is_mature(tip_height) {
-                        immature += txout.txout.value;
-                    }
-                }
-                ObservedAs::Unconfirmed(_) => {
-                    if should_trust(&txout.txout.script_pubkey) {
-                        trusted_pending += txout.txout.value;
-                    } else {
-                        untrusted_pending += txout.txout.value;
-                    }
-                }
-            }
-        }
-
-        Ok(Balance {
-            immature,
-            trusted_pending,
-            untrusted_pending,
-            confirmed,
-        })
-    }
-
-    pub fn balance<C, F>(&self, chain: &C, chain_tip: BlockId, should_trust: F) -> Balance
-    where
-        C: ChainOracle<Error = Infallible>,
-        F: FnMut(&Script) -> bool,
-    {
-        self.try_balance(chain, chain_tip, should_trust)
-            .expect("error is infallible")
-    }
-}
-
 /// A structure that represents changes to an [`IndexedTxGraph`].
 #[derive(Clone, Debug, PartialEq)]
 #[cfg_attr(
@@ -324,9 +220,3 @@ pub trait Indexer {
     /// Determines whether the transaction should be included in the index.
     fn is_tx_relevant(&self, tx: &Transaction) -> bool;
 }
-
-/// A trait that extends [`Indexer`] to also index "owned" script pubkeys.
-pub trait OwnedIndexer: Indexer {
-    /// Determines whether a given script pubkey (`spk`) is owned.
-    fn is_spk_owned(&self, spk: &Script) -> bool;
-}

crates/chain/src/keychain/txout_index.rs

@@ -1,6 +1,6 @@
 use crate::{
     collections::*,
-    indexed_tx_graph::{Indexer, OwnedIndexer},
+    indexed_tx_graph::Indexer,
     miniscript::{Descriptor, DescriptorPublicKey},
     spk_iter::BIP32_MAX_INDEX,
     ForEachTxOut, SpkIterator, SpkTxOutIndex,
@@ -109,12 +109,6 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
     }
 }
 
-impl<K: Clone + Ord + Debug> OwnedIndexer for KeychainTxOutIndex<K> {
-    fn is_spk_owned(&self, spk: &Script) -> bool {
-        self.index_of_spk(spk).is_some()
-    }
-}
-
 impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// Scans an object for relevant outpoints, which are stored and indexed internally.
     ///
@@ -153,6 +147,11 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         &self.inner
     }
 
+    /// Get a reference to the set of indexed outpoints.
+    pub fn outpoints(&self) -> &BTreeSet<((K, u32), OutPoint)> {
+        self.inner.outpoints()
+    }
+
     /// Return a reference to the internal map of the keychain to descriptors.
     pub fn keychains(&self) -> &BTreeMap<K, Descriptor<DescriptorPublicKey>> {
         &self.keychains

crates/chain/src/spk_txout_index.rs

@@ -2,7 +2,7 @@ use core::ops::RangeBounds;
 
 use crate::{
     collections::{hash_map::Entry, BTreeMap, BTreeSet, HashMap},
-    indexed_tx_graph::{Indexer, OwnedIndexer},
+    indexed_tx_graph::Indexer,
     ForEachTxOut,
 };
 use bitcoin::{self, OutPoint, Script, Transaction, TxOut, Txid};
@@ -75,12 +75,6 @@ impl<I: Clone + Ord> Indexer for SpkTxOutIndex<I> {
     }
 }
 
-impl<I: Clone + Ord + 'static> OwnedIndexer for SpkTxOutIndex<I> {
-    fn is_spk_owned(&self, spk: &Script) -> bool {
-        self.spk_indices.get(spk).is_some()
-    }
-}
-
 /// This macro is used instead of a member function of `SpkTxOutIndex`, which would result in a
 /// compiler error[E0521]: "borrowed data escapes out of closure" when we attempt to take a
 /// reference out of the `ForEachTxOut` closure during scanning.
@@ -126,6 +120,11 @@ impl<I: Clone + Ord> SpkTxOutIndex<I> {
         scan_txout!(self, op, txout)
     }
 
+    /// Get a reference to the set of indexed outpoints.
+    pub fn outpoints(&self) -> &BTreeSet<(I, OutPoint)> {
+        &self.spk_txouts
+    }
+
     /// Iterate over all known txouts that spend to tracked script pubkeys.
     pub fn txouts(
         &self,