Persist cached SPKs on wallet creation

Author: j-kon

src/wallet/mod.rs

@@ -355,15 +355,15 @@ impl Wallet {
             ..Default::default()
         };
 
-        let tx_graph = make_indexed_graph(
-            &mut stage,
-            Default::default(),
+        let index = make_keychain_index(
             Default::default(),
             descriptor,
             change_descriptor,
             params.lookahead,
             params.use_spk_cache,
         )?;
+        stage.indexer.merge(index.initial_changeset());
+        let tx_graph = IndexedTxGraph::new(index);
 
         Ok(Wallet {
             signers,
@@ -559,10 +559,9 @@ impl Wallet {
             .map(|(op, _)| op)
             .collect();
 
-        let mut stage = ChangeSet::default();
+        let stage = ChangeSet::default();
 
-        let tx_graph = make_indexed_graph(
-            &mut stage,
+        let tx_graph = make_loaded_indexed_graph(
             changeset.tx_graph,
             changeset.indexer,
             descriptor,
@@ -2828,8 +2827,50 @@ fn new_local_utxo(
     }
 }
 
-fn make_indexed_graph(
-    stage: &mut ChangeSet,
+fn make_keychain_index(
+    indexer_changeset: chain::keychain_txout::ChangeSet,
+    descriptor: ExtendedDescriptor,
+    change_descriptor: Option<ExtendedDescriptor>,
+    lookahead: u32,
+    use_spk_cache: bool,
+) -> Result<KeychainTxOutIndex<KeychainKind>, DescriptorError> {
+    let mut idx = KeychainTxOutIndex::from_changeset(lookahead, use_spk_cache, indexer_changeset);
+
+    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),
+    };
+
+    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")
+                }
+            }
+        })?;
+    assert!(
+        change_descriptor_inserted,
+        "this must be the first time we are seeing this descriptor"
+    );
+
+    Ok(idx)
+}
+
+fn make_loaded_indexed_graph(
     tx_graph_changeset: chain::tx_graph::ChangeSet<ConfirmationBlockTime>,
     indexer_changeset: chain::keychain_txout::ChangeSet,
     descriptor: ExtendedDescriptor,
@@ -2838,50 +2879,25 @@ fn make_indexed_graph(
     use_spk_cache: bool,
 ) -> Result<IndexedTxGraph<ConfirmationBlockTime, KeychainTxOutIndex<KeychainKind>>, DescriptorError>
 {
-    let (indexed_graph, changeset) = IndexedTxGraph::from_changeset(
+    let (indexed_graph, reindex_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),
-            };
-
-            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")
-                        }
-                    }
-                })?;
-            assert!(
-                change_descriptor_inserted,
-                "this must be the first time we are seeing this descriptor"
-            );
-
-            Ok(idx)
+        |idx_cs| {
+            make_keychain_index(
+                idx_cs,
+                descriptor,
+                change_descriptor,
+                lookahead,
+                use_spk_cache,
+            )
         },
     )?;
-    stage.tx_graph.merge(changeset.tx_graph);
-    stage.indexer.merge(changeset.indexer);
+    debug_assert!(
+        reindex_changeset.is_empty(),
+        "loading a wallet should not stage new changes"
+    );
     Ok(indexed_graph)
 }
 

tests/persisted_wallet.rs

@@ -101,22 +101,39 @@ fn wallet_is_persisted() -> anyhow::Result<()> {
         let (external_desc, internal_desc) = get_test_tr_single_sig_xprv_and_change_desc();
 
         // create new wallet
-        let wallet_spk_index = {
+        {
             let mut db = create_db(&file_path)?;
-            let mut wallet = Wallet::create(external_desc, internal_desc)
+            let wallet = Wallet::create(external_desc, internal_desc)
                 .network(Network::Testnet)
                 .use_spk_cache(true)
                 .create_wallet(&mut db)?;
 
-            wallet.reveal_next_address(KeychainKind::External);
+            assert!(
+                wallet.staged().is_none(),
+                "persisted wallet creation should write the initial changeset"
+            );
+        }
+
+        // reload immediately to ensure the initial cached SPKs were persisted at creation time
+        let wallet_spk_index = {
+            let mut db = open_db(&file_path)?;
+            let mut wallet = Wallet::load()
+                .check_network(Network::Testnet)
+                .use_spk_cache(true)
+                .load_wallet(&mut db)?
+                .expect("wallet must exist");
+
+            assert!(wallet.staged().is_none());
+
+            let revealed_external_addr = 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",
+                [(
+                    KeychainKind::External,
+                    [revealed_external_addr.index + DEFAULT_LOOKAHEAD],
+                )],
+                "initial cached SPKs should already be persisted at wallet creation",
             );
 
             // persist new wallet changes