Merge #288: Fix KeychainTxOutIndex recovery logic when spk-cache is enabled

e2814e8 fix: `KeychainTxOutIndex` recovery logic when spk-cache is enabled (志宇)

Pull request description:

  ### Description

  We should ensure we populate the cache (from the changeset) before doing any other operation on `KeychainTxOutIndex` (i.e. inserting descriptors).

  Otherwise we will end up re-deriving spks which have already been cached.

  The test is also corrected to check that we don't duplicate deriving spks. A helper function `check_cache_cs` is added to make the tests more thorough and easier to read.

  ### Changelog notice

  ```md
  Fixed:
  * Recovering from spk-cache now works properly.
  ```

  ### 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 `just p` before pushing

  #### Bugfixes:

  * [x] I've added tests to reproduce the issue which are now passing

  > ~* [ ] This pull request breaks the existing API~
  > ~* [ ] I'm linking the issue being fixed by this PR~

ACKs for top commit:
  ValuedMammal:
    ACK e2814e8
  oleonardolima:
    cACK e2814e8

Tree-SHA512: fdfb2880e90664660634952fd461d01425dcaed38db46b23611d8768cb21924705c0455c3251a2353f5eb6155a302cb7a03a40c27ec8e9221ab4200b20a35f0d

Author: ValuedMammal

wallet/src/wallet/mod.rs

@@ -471,27 +471,24 @@ impl Wallet {
             None => (None, Arc::new(SignersContainer::new())),
         };
 
-        let index = create_indexer(
+        let mut stage = ChangeSet {
+            descriptor: Some(descriptor.clone()),
+            change_descriptor: change_descriptor.clone(),
+            local_chain: chain_changeset,
+            network: Some(network),
+            ..Default::default()
+        };
+
+        let indexed_graph = make_indexed_graph(
+            &mut stage,
+            Default::default(),
+            Default::default(),
             descriptor,
             change_descriptor,
             params.lookahead,
             params.use_spk_cache,
         )?;
 
-        let descriptor = index.get_descriptor(KeychainKind::External).cloned();
-        let change_descriptor = index.get_descriptor(KeychainKind::Internal).cloned();
-        let indexed_graph = IndexedTxGraph::new(index);
-        let indexed_graph_changeset = indexed_graph.initial_changeset();
-
-        let stage = ChangeSet {
-            descriptor,
-            change_descriptor,
-            local_chain: chain_changeset,
-            tx_graph: indexed_graph_changeset.tx_graph,
-            indexer: indexed_graph_changeset.indexer,
-            network: Some(network),
-        };
-
         Ok(Wallet {
             signers,
             change_signers,
@@ -675,20 +672,19 @@ impl Wallet {
             None => Arc::new(SignersContainer::new()),
         };
 
-        let index = create_indexer(
+        let mut stage = ChangeSet::default();
+
+        let indexed_graph = make_indexed_graph(
+            &mut stage,
+            changeset.tx_graph,
+            changeset.indexer,
             descriptor,
             change_descriptor,
             params.lookahead,
             params.use_spk_cache,
         )
         .map_err(LoadError::Descriptor)?;
 
-        let mut indexed_graph = IndexedTxGraph::new(index);
-        indexed_graph.apply_changeset(changeset.indexer.into());
-        indexed_graph.apply_changeset(changeset.tx_graph.into());
-
-        let stage = ChangeSet::default();
-
         Ok(Some(Wallet {
             signers,
             change_signers,
@@ -2687,35 +2683,61 @@ fn new_local_utxo(
     }
 }
 
-fn create_indexer(
+fn make_indexed_graph(
+    stage: &mut ChangeSet,
+    tx_graph_changeset: chain::tx_graph::ChangeSet<ConfirmationBlockTime>,
+    indexer_changeset: chain::keychain_txout::ChangeSet,
     descriptor: ExtendedDescriptor,
     change_descriptor: Option<ExtendedDescriptor>,
     lookahead: u32,
     use_spk_cache: bool,
-) -> Result<KeychainTxOutIndex<KeychainKind>, DescriptorError> {
-    let mut indexer = KeychainTxOutIndex::<KeychainKind>::new(lookahead, use_spk_cache);
-
-    assert!(indexer
-        .insert_descriptor(KeychainKind::External, descriptor)
-        .expect("first descriptor introduced must succeed"));
+) -> Result<IndexedTxGraph<ConfirmationBlockTime, KeychainTxOutIndex<KeychainKind>>, DescriptorError>
+{
+    let (indexed_graph, changeset) = IndexedTxGraph::from_changeset(
+        chain::indexed_tx_graph::ChangeSet {
+            tx_graph: tx_graph_changeset,
+            indexer: indexer_changeset,
+        },
+        |idx_cs| -> Result<KeychainTxOutIndex<KeychainKind>, DescriptorError> {
+            let mut idx = KeychainTxOutIndex::from_changeset(lookahead, use_spk_cache, idx_cs);
+
+            let descriptor_inserted = idx
+                .insert_descriptor(KeychainKind::External, descriptor)
+                .expect("already checked to be a unique, wildcard, non-multipath descriptor");
+            assert!(
+                descriptor_inserted,
+                "this must be the first time we are seeing this descriptor"
+            );
+
+            let change_descriptor = match change_descriptor {
+                Some(change_descriptor) => change_descriptor,
+                None => return Ok(idx),
+            };
 
-    if let Some(change_descriptor) = change_descriptor {
-        assert!(indexer
-            .insert_descriptor(KeychainKind::Internal, change_descriptor)
-            .map_err(|e| {
-                use bdk_chain::indexer::keychain_txout::InsertDescriptorError;
-                match e {
-                    InsertDescriptorError::DescriptorAlreadyAssigned { .. } => {
-                        crate::descriptor::error::Error::ExternalAndInternalAreTheSame
-                    }
-                    InsertDescriptorError::KeychainAlreadyAssigned { .. } => {
-                        unreachable!("this is the first time we're assigning internal")
+            let change_descriptor_inserted = idx
+                .insert_descriptor(KeychainKind::Internal, change_descriptor)
+                .map_err(|e| {
+                    use bdk_chain::indexer::keychain_txout::InsertDescriptorError;
+                    match e {
+                        InsertDescriptorError::DescriptorAlreadyAssigned { .. } => {
+                            crate::descriptor::error::Error::ExternalAndInternalAreTheSame
+                        }
+                        InsertDescriptorError::KeychainAlreadyAssigned { .. } => {
+                            unreachable!("this is the first time we're assigning internal")
+                        }
                     }
-                }
-            })?);
-    }
-
-    Ok(indexer)
+                })?;
+            assert!(
+                change_descriptor_inserted,
+                "this must be the first time we are seeing this descriptor"
+            );
+
+            Ok(idx)
+        },
+    )?;
+    stage.tx_graph.merge(changeset.tx_graph);
+    stage.indexer.merge(changeset.indexer);
+    Ok(indexed_graph)
 }
 
 /// Transforms a [`FeeRate`] to `f64` with unit as sat/vb.

wallet/tests/persisted_wallet.rs

@@ -1,8 +1,9 @@
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::path::Path;
 
 use anyhow::Context;
 use assert_matches::assert_matches;
+use bdk_chain::DescriptorId;
 use bdk_chain::{
     keychain_txout::DEFAULT_LOOKAHEAD, ChainPosition, ConfirmationBlockTime, DescriptorExt,
 };
@@ -14,7 +15,9 @@ use bdk_wallet::{
 use bitcoin::constants::ChainHash;
 use bitcoin::hashes::Hash;
 use bitcoin::key::Secp256k1;
-use bitcoin::{absolute, transaction, Amount, BlockHash, Network, ScriptBuf, Transaction, TxOut};
+use bitcoin::{
+    absolute, secp256k1, transaction, Amount, BlockHash, Network, ScriptBuf, Transaction, TxOut,
+};
 use miniscript::{Descriptor, DescriptorPublicKey};
 
 mod common;
@@ -24,6 +27,57 @@ const DB_MAGIC: &[u8] = &[0x21, 0x24, 0x48];
 
 #[test]
 fn wallet_is_persisted() -> anyhow::Result<()> {
+    type SpkCacheChangeSet = BTreeMap<DescriptorId, BTreeMap<u32, ScriptBuf>>;
+
+    /// Check whether the spk-cache field of the changeset contains the expected spk indices.
+    fn check_cache_cs(
+        cache_cs: &SpkCacheChangeSet,
+        expected: impl IntoIterator<Item = (KeychainKind, impl IntoIterator<Item = u32>)>,
+        msg: impl AsRef<str>,
+    ) {
+        let secp = secp256k1::Secp256k1::new();
+        let (external, internal) = get_test_tr_single_sig_xprv_and_change_desc();
+        let (external_desc, _) = external
+            .into_wallet_descriptor(&secp, Network::Testnet)
+            .unwrap();
+        let (internal_desc, _) = internal
+            .into_wallet_descriptor(&secp, Network::Testnet)
+            .unwrap();
+        let external_did = external_desc.descriptor_id();
+        let internal_did = internal_desc.descriptor_id();
+
+        let cache_cmp = cache_cs
+            .iter()
+            .map(|(did, spks)| {
+                let kind: KeychainKind;
+                if did == &external_did {
+                    kind = KeychainKind::External;
+                } else if did == &internal_did {
+                    kind = KeychainKind::Internal;
+                } else {
+                    unreachable!();
+                }
+                let spk_indices = spks.keys().copied().collect::<BTreeSet<u32>>();
+                (kind, spk_indices)
+            })
+            .filter(|(_, spk_indices)| !spk_indices.is_empty())
+            .collect::<BTreeMap<KeychainKind, BTreeSet<u32>>>();
+
+        let expected_cmp = expected
+            .into_iter()
+            .map(|(kind, indices)| (kind, indices.into_iter().collect::<BTreeSet<u32>>()))
+            .filter(|(_, spk_indices)| !spk_indices.is_empty())
+            .collect::<BTreeMap<KeychainKind, BTreeSet<u32>>>();
+
+        assert_eq!(cache_cmp, expected_cmp, "{}", msg.as_ref());
+    }
+
+    fn staged_cache(wallet: &Wallet) -> SpkCacheChangeSet {
+        wallet.staged().map_or(SpkCacheChangeSet::default(), |cs| {
+            cs.indexer.spk_cache.clone()
+        })
+    }
+
     fn run<Db, CreateDb, OpenDb>(
         filename: &str,
         create_db: CreateDb,
@@ -46,8 +100,18 @@ fn wallet_is_persisted() -> anyhow::Result<()> {
                 .network(Network::Testnet)
                 .use_spk_cache(true)
                 .create_wallet(&mut db)?;
+
             wallet.reveal_next_address(KeychainKind::External);
 
+            check_cache_cs(
+                &staged_cache(&wallet),
+                [
+                    (KeychainKind::External, 0..DEFAULT_LOOKAHEAD + 1),
+                    (KeychainKind::Internal, 0..DEFAULT_LOOKAHEAD),
+                ],
+                "cache cs must return initial set + the external index that was just derived",
+            );
+
             // persist new wallet changes
             assert!(wallet.persist(&mut db)?, "must write");
             wallet.spk_index().clone()
@@ -81,6 +145,7 @@ fn wallet_is_persisted() -> anyhow::Result<()> {
                     .0
             );
         }
+
         // Test SPK cache
         {
             let mut db = open_db(&file_path).context("failed to recover db")?;
@@ -90,35 +155,32 @@ fn wallet_is_persisted() -> anyhow::Result<()> {
                 .load_wallet(&mut db)?
                 .expect("wallet must exist");
 
-            let external_did = wallet
-                .public_descriptor(KeychainKind::External)
-                .descriptor_id();
-            let internal_did = wallet
-                .public_descriptor(KeychainKind::Internal)
-                .descriptor_id();
-
             assert!(wallet.staged().is_none());
 
-            let _addr = wallet.reveal_next_address(KeychainKind::External);
-            let cs = wallet.staged().expect("we should have staged a changeset");
-            assert!(!cs.indexer.spk_cache.is_empty(), "failed to cache spks");
-            assert_eq!(cs.indexer.spk_cache.len(), 2, "we persisted two keychains");
-            let spk_cache: &BTreeMap<u32, ScriptBuf> =
-                cs.indexer.spk_cache.get(&external_did).unwrap();
-            assert_eq!(spk_cache.len() as u32, 1 + 1 + DEFAULT_LOOKAHEAD);
-            assert_eq!(spk_cache.keys().last(), Some(&26));
-            let spk_cache = cs.indexer.spk_cache.get(&internal_did).unwrap();
-            assert_eq!(spk_cache.len() as u32, DEFAULT_LOOKAHEAD);
-            assert_eq!(spk_cache.keys().last(), Some(&24));
+            let revealed_external_addr = wallet.reveal_next_address(KeychainKind::External);
+            check_cache_cs(
+                &staged_cache(&wallet),
+                [(
+                    KeychainKind::External,
+                    [revealed_external_addr.index + DEFAULT_LOOKAHEAD],
+                )],
+                "must only persist the revealed+LOOKAHEAD indexed external spk",
+            );
+
             // Clear the stage
             let _ = wallet.take_staged();
-            let _addr = wallet.reveal_next_address(KeychainKind::Internal);
-            let cs = wallet.staged().unwrap();
-            assert_eq!(cs.indexer.spk_cache.len(), 1);
-            let spk_cache = cs.indexer.spk_cache.get(&internal_did).unwrap();
-            assert_eq!(spk_cache.len(), 1);
-            assert_eq!(spk_cache.keys().next(), Some(&25));
+
+            let revealed_internal_addr = wallet.reveal_next_address(KeychainKind::Internal);
+            check_cache_cs(
+                &staged_cache(&wallet),
+                [(
+                    KeychainKind::Internal,
+                    [revealed_internal_addr.index + DEFAULT_LOOKAHEAD],
+                )],
+                "must only persist the revealed+LOOKAHEAD indexed internal spk",
+            );
         }
+
         // SPK cache requires load params
         {
             let mut db = open_db(&file_path).context("failed to recover db")?;