feat(chain): Add TxGraph methods that handle expected spk txids

* `TxGraph::try_list_expected_spk_txids`
* `TxGraph::list_expected_spk_txids`
* `IndexedTxGraph::try_list_expected_spk_txids`
* `IndexedTxGraph::list_expected_spk_txids`

Co-authored-by: valued mammal <[email protected]>
Co-authored-by: Wei Chen <[email protected]>

Author: 3 people

crates/chain/src/indexed_tx_graph.rs

@@ -1,13 +1,18 @@
 //! Contains the [`IndexedTxGraph`] and associated types. Refer to the
 //! [`IndexedTxGraph`] documentation for more.
-use core::fmt::Debug;
+use core::{
+    convert::Infallible,
+    fmt::{self, Debug},
+    ops::RangeBounds,
+};
 
 use alloc::{sync::Arc, vec::Vec};
-use bitcoin::{Block, OutPoint, Transaction, TxOut, Txid};
+use bitcoin::{Block, OutPoint, ScriptBuf, Transaction, TxOut, Txid};
 
 use crate::{
+    spk_txout::SpkTxOutIndex,
     tx_graph::{self, TxGraph},
-    Anchor, BlockId, Indexer, Merge, TxPosInBlock,
+    Anchor, BlockId, ChainOracle, Indexer, Merge, TxPosInBlock,
 };
 
 /// The [`IndexedTxGraph`] combines a [`TxGraph`] and an [`Indexer`] implementation.
@@ -127,6 +132,19 @@ where
         self.graph.insert_seen_at(txid, seen_at).into()
     }
 
+    /// Inserts the given `evicted_at` for `txid`.
+    ///
+    /// The `evicted_at` timestamp represents the last known time when the transaction was observed
+    /// to be missing from the mempool. If `txid` was previously recorded with an earlier
+    /// `evicted_at` value, it is updated only if the new value is greater.
+    pub fn insert_evicted_at(&mut self, txid: Txid, evicted_at: u64) -> ChangeSet<A, I::ChangeSet> {
+        let tx_graph = self.graph.insert_evicted_at(txid, evicted_at);
+        ChangeSet {
+            tx_graph,
+            ..Default::default()
+        }
+    }
+
     /// Batch insert transactions, filtering out those that are irrelevant.
     ///
     /// Relevancy is determined by the [`Indexer::is_tx_relevant`] implementation of `I`. Irrelevant
@@ -301,6 +319,58 @@ where
     }
 }
 
+impl<A, X> IndexedTxGraph<A, X>
+where
+    A: Anchor,
+{
+    /// List txids that are expected to exist under the given spks.
+    ///
+    /// This is used to fill [`SyncRequestBuilder::expected_spk_txids`](bdk_core::spk_client::SyncRequestBuilder::expected_spk_txids).
+    ///
+    /// The spk index range can be contrained with `range`.
+    ///
+    /// # Error
+    ///
+    /// If the [`ChainOracle`] implementation (`chain`) fails, an error will be returned with the
+    /// returned item.
+    ///
+    /// If the [`ChainOracle`] is infallible,
+    /// [`list_expected_spk_txids`](Self::list_expected_spk_txids) can be used instead.
+    pub fn try_list_expected_spk_txids<'a, C, I>(
+        &'a self,
+        chain: &'a C,
+        chain_tip: BlockId,
+        spk_index_range: impl RangeBounds<I> + 'a,
+    ) -> impl Iterator<Item = Result<(ScriptBuf, Txid), C::Error>> + 'a
+    where
+        C: ChainOracle,
+        X: AsRef<SpkTxOutIndex<I>> + 'a,
+        I: fmt::Debug + Clone + Ord + 'a,
+    {
+        self.graph
+            .try_list_expected_spk_txids(chain, chain_tip, &self.index, spk_index_range)
+    }
+
+    /// List txids that are expected to exist under the given spks.
+    ///
+    /// This is the infallible version of
+    /// [`try_list_expected_spk_txids`](Self::try_list_expected_spk_txids).
+    pub fn list_expected_spk_txids<'a, C, I>(
+        &'a self,
+        chain: &'a C,
+        chain_tip: BlockId,
+        spk_index_range: impl RangeBounds<I> + 'a,
+    ) -> impl Iterator<Item = (ScriptBuf, Txid)> + 'a
+    where
+        C: ChainOracle<Error = Infallible>,
+        X: AsRef<SpkTxOutIndex<I>> + 'a,
+        I: fmt::Debug + Clone + Ord + 'a,
+    {
+        self.try_list_expected_spk_txids(chain, chain_tip, spk_index_range)
+            .map(|r| r.expect("infallible"))
+    }
+}
+
 impl<A, I> AsRef<TxGraph<A>> for IndexedTxGraph<A, I> {
     fn as_ref(&self) -> &TxGraph<A> {
         &self.graph

crates/chain/src/indexer/keychain_txout.rs

@@ -136,6 +136,12 @@ impl<K> Default for KeychainTxOutIndex<K> {
     }
 }
 
+impl<K> AsRef<SpkTxOutIndex<(K, u32)>> for KeychainTxOutIndex<K> {
+    fn as_ref(&self) -> &SpkTxOutIndex<(K, u32)> {
+        &self.inner
+    }
+}
+
 impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
     type ChangeSet = ChangeSet;
 
@@ -200,6 +206,11 @@ impl<K> KeychainTxOutIndex<K> {
             lookahead,
         }
     }
+
+    /// Get a reference to the internal [`SpkTxOutIndex`].
+    pub fn inner(&self) -> &SpkTxOutIndex<(K, u32)> {
+        &self.inner
+    }
 }
 
 /// Methods that are *re-exposed* from the internal [`SpkTxOutIndex`].

crates/chain/src/indexer/spk_txout.rs

@@ -54,6 +54,12 @@ impl<I> Default for SpkTxOutIndex<I> {
     }
 }
 
+impl<I> AsRef<SpkTxOutIndex<I>> for SpkTxOutIndex<I> {
+    fn as_ref(&self) -> &SpkTxOutIndex<I> {
+        self
+    }
+}
+
 impl<I: Clone + Ord + core::fmt::Debug> Indexer for SpkTxOutIndex<I> {
     type ChangeSet = ();
 
@@ -334,4 +340,24 @@ impl<I: Clone + Ord + core::fmt::Debug> SpkTxOutIndex<I> {
             .any(|output| self.spk_indices.contains_key(&output.script_pubkey));
         input_matches || output_matches
     }
+
+    /// Find relevant script pubkeys associated with a transaction for tracking and validation.
+    ///
+    /// Returns a set of script pubkeys from [`SpkTxOutIndex`] that are relevant to the outputs and
+    /// previous outputs of a given transaction. Inputs are only considered relevant if the parent
+    /// transactions have been scanned.
+    pub fn relevant_spks_of_tx(&self, tx: &Transaction) -> BTreeSet<(I, ScriptBuf)> {
+        let spks_from_inputs = tx.input.iter().filter_map(|txin| {
+            self.txouts
+                .get(&txin.previous_output)
+                .cloned()
+                .map(|(i, prev_txo)| (i, prev_txo.script_pubkey))
+        });
+        let spks_from_outputs = tx
+            .output
+            .iter()
+            .filter_map(|txout| self.spk_indices.get_key_value(&txout.script_pubkey))
+            .map(|(spk, i)| (i.clone(), spk.clone()));
+        spks_from_inputs.chain(spks_from_outputs).collect()
+    }
 }

crates/chain/src/tx_graph.rs

@@ -120,6 +120,7 @@
 //! [`insert_txout`]: TxGraph::insert_txout
 
 use crate::collections::*;
+use crate::spk_txout::SpkTxOutIndex;
 use crate::BlockId;
 use crate::CanonicalIter;
 use crate::CanonicalReason;
@@ -132,6 +133,7 @@ use bdk_core::ConfirmationBlockTime;
 pub use bdk_core::TxUpdate;
 use bitcoin::{Amount, OutPoint, ScriptBuf, SignedAmount, Transaction, TxOut, Txid};
 use core::fmt::{self, Formatter};
+use core::ops::RangeBounds;
 use core::{
     convert::Infallible,
     ops::{Deref, RangeInclusive},
@@ -1150,6 +1152,67 @@ impl<A: Anchor> TxGraph<A> {
         self.try_balance(chain, chain_tip, outpoints, trust_predicate)
             .expect("oracle is infallible")
     }
+
+    /// List txids that are expected to exist under the given spks.
+    ///
+    /// This is used to fill [`SyncRequestBuilder::expected_spk_txids`](bdk_core::spk_client::SyncRequestBuilder::expected_spk_txids).
+    ///
+    /// The spk index range can be contrained with `range`.
+    ///
+    /// # Error
+    ///
+    /// If the [`ChainOracle`] implementation (`chain`) fails, an error will be returned with the
+    /// returned item.
+    ///
+    /// If the [`ChainOracle`] is infallible,
+    /// [`list_expected_spk_txids`](Self::list_expected_spk_txids) can be used instead.
+    pub fn try_list_expected_spk_txids<'a, C, I>(
+        &'a self,
+        chain: &'a C,
+        chain_tip: BlockId,
+        indexer: &'a impl AsRef<SpkTxOutIndex<I>>,
+        spk_index_range: impl RangeBounds<I> + 'a,
+    ) -> impl Iterator<Item = Result<(ScriptBuf, Txid), C::Error>> + 'a
+    where
+        C: ChainOracle,
+        I: fmt::Debug + Clone + Ord + 'a,
+    {
+        let indexer = indexer.as_ref();
+        self.try_list_canonical_txs(chain, chain_tip).flat_map(
+            move |res| -> Vec<Result<(ScriptBuf, Txid), C::Error>> {
+                let range = &spk_index_range;
+                let c_tx = match res {
+                    Ok(c_tx) => c_tx,
+                    Err(err) => return vec![Err(err)],
+                };
+                let relevant_spks = indexer.relevant_spks_of_tx(&c_tx.tx_node);
+                relevant_spks
+                    .into_iter()
+                    .filter(|(i, _)| range.contains(i))
+                    .map(|(_, spk)| Ok((spk, c_tx.tx_node.txid)))
+                    .collect()
+            },
+        )
+    }
+
+    /// List txids that are expected to exist under the given spks.
+    ///
+    /// This is the infallible version of
+    /// [`try_list_expected_spk_txids`](Self::try_list_expected_spk_txids).
+    pub fn list_expected_spk_txids<'a, C, I>(
+        &'a self,
+        chain: &'a C,
+        chain_tip: BlockId,
+        indexer: &'a impl AsRef<SpkTxOutIndex<I>>,
+        spk_index_range: impl RangeBounds<I> + 'a,
+    ) -> impl Iterator<Item = (ScriptBuf, Txid)> + 'a
+    where
+        C: ChainOracle<Error = Infallible>,
+        I: fmt::Debug + Clone + Ord + 'a,
+    {
+        self.try_list_expected_spk_txids(chain, chain_tip, indexer, spk_index_range)
+            .map(|r| r.expect("infallible"))
+    }
 }
 
 /// The [`ChangeSet`] represents changes to a [`TxGraph`].