Merge #1968: bench: Add reindex_tx_graph benchmark

f51f5b5 docs(chain): Improve API docs (valued mammal)
fcf3834 test(keychain_txout): test spk cache (valued mammal)
e07ec1c bench(chain): Add `reindex_tx_graph` benchmark (valued mammal)

Pull request description:

  This PR contains the following changes:

  - Add benchmark `reindex_tx_graph`. Run with `cargo bench -p bdk_chain --bench indexer`.
  - Add unit test to `keychain_txout` module to test behavior of spk-cache.
  - Fixup a few doc comments.

  ### Checklists

  #### All Submissions:

  * [x] I've signed all my commits
  * [x] I followed the [contribution guidelines](https://github.com/bitcoindevkit/bdk/blob/master/CONTRIBUTING.md)
  * [x] I ran `cargo +nightly fmt` and `cargo clippy` before committing

ACKs for top commit:
  ValuedMammal:
    ACK f51f5b5

Tree-SHA512: 4d6826a825fa40fa13bd622ed6cff49366c541a7e0a7035a45fbeb98853d693ddef7fab113f6d77259fd75eddfc91623834363e028235fb26336c8865fccb02b

Author: ValuedMammal

crates/chain/Cargo.toml

@@ -40,3 +40,7 @@ rusqlite = ["std", "dep:rusqlite", "serde"]
 [[bench]]
 name = "canonicalization"
 harness = false
+
+[[bench]]
+name = "indexer"
+harness = false

crates/chain/benches/indexer.rs

@@ -0,0 +1,121 @@
+use bdk_chain::{
+    keychain_txout::{InsertDescriptorError, KeychainTxOutIndex},
+    local_chain::LocalChain,
+    CanonicalizationParams, IndexedTxGraph,
+};
+use bdk_core::{BlockId, CheckPoint, ConfirmationBlockTime, TxUpdate};
+use bitcoin::{
+    absolute, constants, hashes::Hash, key::Secp256k1, transaction, Amount, BlockHash, Network,
+    Transaction, TxIn, TxOut,
+};
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use miniscript::Descriptor;
+use std::sync::Arc;
+
+type Keychain = ();
+type KeychainTxGraph = IndexedTxGraph<ConfirmationBlockTime, KeychainTxOutIndex<Keychain>>;
+
+const DESC: &str = "tr([ab28dc00/86h/1h/0h]tpubDCdDtzAMZZrkwKBxwNcGCqe4FRydeD9rfMisoi7qLdraG79YohRfPW4YgdKQhpgASdvh612xXNY5xYzoqnyCgPbkpK4LSVcH5Xv4cK7johH/0/*)";
+const LOOKAHEAD: u32 = 10;
+const LAST_REVEALED: u32 = 500;
+const TX_CT: u32 = 21;
+const USE_SPK_CACHE: bool = true;
+const AMOUNT: Amount = Amount::from_sat(1_000);
+
+fn new_tx(lt: u32) -> Transaction {
+    Transaction {
+        version: transaction::Version::TWO,
+        lock_time: absolute::LockTime::from_consensus(lt),
+        input: vec![],
+        output: vec![TxOut::NULL],
+    }
+}
+
+fn genesis_block_id() -> BlockId {
+    BlockId {
+        height: 0,
+        hash: constants::genesis_block(Network::Regtest).block_hash(),
+    }
+}
+
+fn tip_block_id() -> BlockId {
+    BlockId {
+        height: 100,
+        hash: BlockHash::all_zeros(),
+    }
+}
+
+fn setup<F: Fn(&mut KeychainTxGraph, &LocalChain)>(f: F) -> (KeychainTxGraph, LocalChain) {
+    let desc = Descriptor::parse_descriptor(&Secp256k1::new(), DESC)
+        .unwrap()
+        .0;
+
+    let cp = CheckPoint::from_block_ids([genesis_block_id(), tip_block_id()]).unwrap();
+    let chain = LocalChain::from_tip(cp).unwrap();
+
+    let mut index = KeychainTxOutIndex::new(LOOKAHEAD, USE_SPK_CACHE);
+    index.insert_descriptor((), desc).unwrap();
+    let mut tx_graph = KeychainTxGraph::new(index);
+
+    f(&mut tx_graph, &chain);
+
+    (tx_graph, chain)
+}
+
+/// Bench performance of recovering `KeychainTxOutIndex` from changeset.
+fn do_bench(indexed_tx_graph: &KeychainTxGraph, chain: &LocalChain) {
+    let desc = indexed_tx_graph.index.get_descriptor(()).unwrap();
+    let changeset = indexed_tx_graph.initial_changeset();
+
+    // Now recover
+    let (graph, _cs) =
+        KeychainTxGraph::from_changeset(changeset, |cs| -> Result<_, InsertDescriptorError<_>> {
+            let mut index = KeychainTxOutIndex::from_changeset(LOOKAHEAD, USE_SPK_CACHE, cs);
+            let _ = index.insert_descriptor((), desc.clone())?;
+            Ok(index)
+        })
+        .unwrap();
+
+    // Check balance
+    let chain_tip = chain.tip().block_id();
+    let op = graph.index.outpoints().clone();
+    let bal = graph.graph().balance(
+        chain,
+        chain_tip,
+        CanonicalizationParams::default(),
+        op,
+        |_, _| false,
+    );
+    assert_eq!(bal.total(), AMOUNT * TX_CT as u64);
+}
+
+pub fn reindex_tx_graph(c: &mut Criterion) {
+    let (graph, chain) = black_box(setup(|graph, _chain| {
+        // Add relevant txs to graph
+        for i in 0..TX_CT {
+            let script_pubkey = graph.index.reveal_next_spk(()).unwrap().0 .1;
+            let tx = Transaction {
+                input: vec![TxIn::default()],
+                output: vec![TxOut {
+                    script_pubkey,
+                    value: AMOUNT,
+                }],
+                ..new_tx(i)
+            };
+            let txid = tx.compute_txid();
+            let mut update = TxUpdate::default();
+            update.seen_ats = [(txid, i as u64)].into();
+            update.txs = vec![Arc::new(tx)];
+            let _ = graph.apply_update(update);
+        }
+        // Reveal some SPKs
+        let _ = graph.index.reveal_to_target((), LAST_REVEALED);
+    }));
+
+    c.bench_function("reindex_tx_graph", {
+        move |b| b.iter(|| do_bench(&graph, &chain))
+    });
+}
+
+criterion_group!(benches, reindex_tx_graph);
+criterion_main!(benches);

crates/chain/src/indexed_tx_graph.rs

@@ -112,7 +112,7 @@ where
     ///
     /// let (graph, reindex_cs) =
     ///     IndexedTxGraph::from_changeset(persisted_changeset, move |idx_cs| -> anyhow::Result<_> {
-    ///         // e.g. KeychainTxOutIndex needs descriptors that weren’t in its CS
+    ///         // e.g. KeychainTxOutIndex needs descriptors that weren’t in its change set.
     ///         let mut idx = KeychainTxOutIndex::from_changeset(DEFAULT_LOOKAHEAD, true, idx_cs);
     ///         if let Some(desc) = persisted_desc {
     ///             idx.insert_descriptor("external", desc)?;

crates/chain/src/indexer/keychain_txout.rs

@@ -200,7 +200,7 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
 }
 
 impl<K> KeychainTxOutIndex<K> {
-    /// Construct a [`KeychainTxOutIndex`] with the given `lookahead` and `use_spk_cache` boolean.
+    /// Construct a [`KeychainTxOutIndex`] with the given `lookahead` and `persist_spks` boolean.
     ///
     /// # Lookahead
     ///
@@ -221,10 +221,10 @@ impl<K> KeychainTxOutIndex<K> {
     ///
     /// ```rust
     /// # use bdk_chain::keychain_txout::KeychainTxOutIndex;
-    /// // Derive 20 future addresses per chain and persist + reload script pubkeys via ChangeSets:
+    /// // Derive 20 future addresses per keychain and persist + reload script pubkeys via ChangeSets:
     /// let idx = KeychainTxOutIndex::<&'static str>::new(20, true);
     ///
-    /// // Derive 10 future addresses per chain without persistence:
+    /// // Derive 10 future addresses per keychain without persistence:
     /// let idx = KeychainTxOutIndex::<&'static str>::new(10, false);
     /// ```
     pub fn new(lookahead: u32, persist_spks: bool) -> Self {
@@ -251,7 +251,7 @@ impl<K> KeychainTxOutIndex<K> {
 impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// Construct `KeychainTxOutIndex<K>` from the given `changeset`.
     ///
-    /// Shorthand for called [`new`] and then [`apply_changeset`].
+    /// Shorthand for calling [`new`] and then [`apply_changeset`].
     ///
     /// [`new`]: Self::new
     /// [`apply_changeset`]: Self::apply_changeset
@@ -1002,7 +1002,7 @@ impl<K: core::fmt::Debug> std::error::Error for InsertDescriptorError<K> {}
 ///
 /// It tracks:
 /// 1. `last_revealed`: the highest derivation index revealed per descriptor.
-/// 2. `spks`: the cache of derived script pubkeys to persist across runs.
+/// 2. `spk_cache`: the cache of derived script pubkeys to persist across runs.
 ///
 /// You can apply a `ChangeSet` to a `KeychainTxOutIndex` via
 /// [`KeychainTxOutIndex::apply_changeset`], or merge two change sets with [`ChangeSet::merge`].
@@ -1011,7 +1011,7 @@ impl<K: core::fmt::Debug> std::error::Error for InsertDescriptorError<K> {}
 ///
 /// - `last_revealed` is monotonic: merging retains the maximum index for each descriptor and never
 ///   decreases.
-/// - `spks` accumulates entries: once a script pubkey is persisted, it remains available for
+/// - `spk_cache` accumulates entries: once a script pubkey is persisted, it remains available for
 ///   reload. If the same descriptor and index appear again with a new script pubkey, the latter
 ///   value overrides the former.
 ///
@@ -1107,3 +1107,70 @@ impl<K: Clone + Ord + core::fmt::Debug> FullScanRequestBuilderExt<K> for FullSca
         self
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    use bdk_testenv::utils::DESCRIPTORS;
+    use bitcoin::secp256k1::Secp256k1;
+    use miniscript::Descriptor;
+
+    // Test that `KeychainTxOutIndex` uses the spk cache.
+    // And the indexed spks are as expected.
+    #[test]
+    fn test_spk_cache() {
+        let lookahead = 10;
+        let use_cache = true;
+        let mut index = KeychainTxOutIndex::new(lookahead, use_cache);
+        let s = DESCRIPTORS[0];
+
+        let desc = Descriptor::parse_descriptor(&Secp256k1::new(), s)
+            .unwrap()
+            .0;
+
+        let did = desc.descriptor_id();
+
+        let reveal_to = 2;
+        let end_index = reveal_to + lookahead;
+
+        let _ = index.insert_descriptor(0i32, desc.clone());
+        assert_eq!(index.spk_cache.get(&did).unwrap().len() as u32, lookahead);
+        assert_eq!(index.next_index(0), Some((0, true)));
+
+        // Now reveal some scripts
+        for _ in 0..=reveal_to {
+            let _ = index.reveal_next_spk(0).unwrap();
+        }
+        assert_eq!(index.last_revealed_index(0), Some(reveal_to));
+
+        let spk_cache = &index.spk_cache;
+        assert!(!spk_cache.is_empty());
+
+        for (&did, cached_spks) in spk_cache {
+            assert_eq!(did, desc.descriptor_id());
+            for (&i, cached_spk) in cached_spks {
+                // Cached spk matches derived
+                let exp_spk = desc.at_derivation_index(i).unwrap().script_pubkey();
+                assert_eq!(&exp_spk, cached_spk);
+                // Also matches the inner index
+                assert_eq!(index.spk_at_index(0, i), Some(cached_spk.clone()));
+            }
+        }
+
+        let init_cs = index.initial_changeset();
+        assert_eq!(
+            init_cs.spk_cache.get(&did).unwrap().len() as u32,
+            end_index + 1
+        );
+
+        // Now test load from changeset
+        let recovered =
+            KeychainTxOutIndex::<&str>::from_changeset(lookahead, use_cache, init_cs.clone());
+        assert_eq!(&recovered.spk_cache, spk_cache);
+
+        // The cache is optional at load time
+        let index = KeychainTxOutIndex::<i32>::from_changeset(lookahead, false, init_cs);
+        assert!(index.spk_cache.is_empty());
+    }
+}