fix(keychain_txout): Remove spk_cache_stage

`spk_cache_stage` was being used to hold newly derived
SPKs before being transferred to the `ChangeSet`. As a result
it became necessary to call `_empty_stage_into_changeset` whenever
the internal state changed, requiring the developer to know
which methods depended on the functionality, e.g. `scan_txout`
and `reveal_next_spk`.

Further, this altered the behavior of the `ChangeSet` in
surprising ways. For example revealing 1 SPK might return
additional scripts in the `ChangeSet` if they had
previously been staged.

The rationale in bitcoindevkit#1963 was to avoid returning a `ChangeSet` from
`insert_descriptor`. We can do this by providing another
method that returns a tuple of `(bool, ChangeSet)` so that
descriptor insertion results in updating the `ChangeSet` directly.
Any code using the original `insert_descriptor` can still access the
`ChangeSet` using `initial_changeset` (which may include SPKs
from other descriptors), or we can expose an accessor to get a
reference to the `spk_cache`.

Author: ValuedMammal

crates/chain/src/indexer/keychain_txout.rs

@@ -138,8 +138,6 @@ pub struct KeychainTxOutIndex<K> {
     persist_spks: bool,
     /// Cache of derived spks.
     spk_cache: BTreeMap<DescriptorId, HashMap<u32, ScriptBuf>>,
-    /// Staged script pubkeys waiting to be written out in the next ChangeSet.
-    spk_cache_stage: BTreeMap<DescriptorId, Vec<(u32, ScriptBuf)>>,
 }
 
 impl<K> Default for KeychainTxOutIndex<K> {
@@ -160,7 +158,6 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
     fn index_txout(&mut self, outpoint: OutPoint, txout: &TxOut) -> Self::ChangeSet {
         let mut changeset = ChangeSet::default();
         self._index_txout(&mut changeset, outpoint, txout);
-        self._empty_stage_into_changeset(&mut changeset);
         changeset
     }
 
@@ -170,7 +167,6 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
         for (vout, txout) in tx.output.iter().enumerate() {
             self._index_txout(&mut changeset, OutPoint::new(txid, vout as u32), txout);
         }
-        self._empty_stage_into_changeset(&mut changeset);
         changeset
     }
 
@@ -237,7 +233,6 @@ impl<K> KeychainTxOutIndex<K> {
             lookahead,
             persist_spks,
             spk_cache: Default::default(),
-            spk_cache_stage: Default::default(),
         }
     }
 
@@ -280,33 +275,11 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             };
             if index_updated {
                 changeset.last_revealed.insert(*did, index);
-                self.replenish_inner_index(*did, &keychain, self.lookahead);
+                self.replenish_inner_index(*did, &keychain, self.lookahead, changeset);
             }
         }
     }
 
-    fn _empty_stage_into_changeset(&mut self, changeset: &mut ChangeSet) {
-        if !self.persist_spks {
-            return;
-        }
-        for (did, spks) in core::mem::take(&mut self.spk_cache_stage) {
-            debug_assert!(
-                {
-                    let desc = self.descriptors.get(&did).expect("invariant");
-                    spks.iter().all(|(i, spk)| {
-                        let exp_spk = desc
-                            .at_derivation_index(*i)
-                            .expect("must derive")
-                            .script_pubkey();
-                        &exp_spk == spk
-                    })
-                },
-                "all staged spks must be correct"
-            );
-            changeset.spk_cache.entry(did).or_default().extend(spks);
-        }
-    }
-
     /// Get the set of indexed outpoints, corresponding to tracked keychains.
     pub fn outpoints(&self) -> &BTreeSet<KeychainIndexed<K, OutPoint>> {
         self.inner.outpoints()
@@ -468,7 +441,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             self.descriptors.insert(did, descriptor.clone());
             self.keychain_to_descriptor_id.insert(keychain.clone(), did);
             self.descriptor_id_to_keychain.insert(did, keychain.clone());
-            self.replenish_inner_index(did, &keychain, self.lookahead);
+            self.replenish_inner_index(did, &keychain, self.lookahead, &mut ChangeSet::default());
             return Ok(true);
         }
 
@@ -523,27 +496,42 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
                 .filter(|&index| index > 0);
 
             if let Some(temp_lookahead) = temp_lookahead {
-                self.replenish_inner_index_keychain(keychain, temp_lookahead);
+                self.replenish_inner_index_keychain(keychain, temp_lookahead, &mut changeset);
             }
         }
-        self._empty_stage_into_changeset(&mut changeset);
         changeset
     }
 
-    fn replenish_inner_index_did(&mut self, did: DescriptorId, lookahead: u32) {
+    fn replenish_inner_index_did(
+        &mut self,
+        did: DescriptorId,
+        lookahead: u32,
+        changeset: &mut ChangeSet,
+    ) {
         if let Some(keychain) = self.descriptor_id_to_keychain.get(&did).cloned() {
-            self.replenish_inner_index(did, &keychain, lookahead);
+            self.replenish_inner_index(did, &keychain, lookahead, changeset);
         }
     }
 
-    fn replenish_inner_index_keychain(&mut self, keychain: K, lookahead: u32) {
+    fn replenish_inner_index_keychain(
+        &mut self,
+        keychain: K,
+        lookahead: u32,
+        changeset: &mut ChangeSet,
+    ) {
         if let Some(did) = self.keychain_to_descriptor_id.get(&keychain) {
-            self.replenish_inner_index(*did, &keychain, lookahead);
+            self.replenish_inner_index(*did, &keychain, lookahead, changeset);
         }
     }
 
     /// Syncs the state of the inner spk index after changes to a keychain
-    fn replenish_inner_index(&mut self, did: DescriptorId, keychain: &K, lookahead: u32) {
+    fn replenish_inner_index(
+        &mut self,
+        did: DescriptorId,
+        keychain: &K,
+        lookahead: u32,
+        changeset: &mut ChangeSet,
+    ) {
         let descriptor = self.descriptors.get(&did).expect("invariant");
 
         let mut next_index = self
@@ -574,7 +562,6 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             };
             let cached_spk_iter = core::iter::from_fn({
                 let spk_cache = self.spk_cache.entry(did).or_default();
-                let spk_stage = self.spk_cache_stage.entry(did).or_default();
                 let _i = &mut next_index;
                 move || -> Option<Indexed<ScriptBuf>> {
                     if *_i >= stop_index {
@@ -588,8 +575,12 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
                         return Some((spk_i, spk.clone()));
                     }
                     let spk = derive_spk(spk_i);
-                    spk_stage.push((spk_i, spk.clone()));
                     spk_cache.insert(spk_i, spk.clone());
+                    changeset
+                        .spk_cache
+                        .entry(did)
+                        .or_default()
+                        .insert(spk_i, spk.clone());
                     Some((spk_i, spk))
                 }
             });
@@ -778,7 +769,6 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             self._reveal_to_target(&mut changeset, keychain.clone(), index);
         }
 
-        self._empty_stage_into_changeset(&mut changeset);
         changeset
     }
 
@@ -802,7 +792,6 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     ) -> Option<(Vec<Indexed<ScriptBuf>>, ChangeSet)> {
         let mut changeset = ChangeSet::default();
         let revealed_spks = self._reveal_to_target(&mut changeset, keychain, target_index)?;
-        self._empty_stage_into_changeset(&mut changeset);
         Some((revealed_spks, changeset))
     }
     fn _reveal_to_target(
@@ -840,7 +829,6 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     pub fn reveal_next_spk(&mut self, keychain: K) -> Option<(Indexed<ScriptBuf>, ChangeSet)> {
         let mut changeset = ChangeSet::default();
         let indexed_spk = self._reveal_next_spk(&mut changeset, keychain)?;
-        self._empty_stage_into_changeset(&mut changeset);
         Some((indexed_spk, changeset))
     }
     fn _reveal_next_spk(
@@ -853,7 +841,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             let did = self.keychain_to_descriptor_id.get(&keychain)?;
             self.last_revealed.insert(*did, next_index);
             changeset.last_revealed.insert(*did, next_index);
-            self.replenish_inner_index(*did, &keychain, self.lookahead);
+            self.replenish_inner_index(*did, &keychain, self.lookahead, changeset);
         }
         let script = self
             .inner
@@ -882,7 +870,6 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
             .next()
             .map(|(i, spk)| (i, spk.to_owned()));
         let spk = next_unused.or_else(|| self._reveal_next_spk(&mut changeset, keychain))?;
-        self._empty_stage_into_changeset(&mut changeset);
         Some((spk, changeset))
     }
 
@@ -949,7 +936,7 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
         for (did, index) in changeset.last_revealed {
             let v = self.last_revealed.entry(did).or_default();
             *v = index.max(*v);
-            self.replenish_inner_index_did(did, self.lookahead);
+            self.replenish_inner_index_did(did, self.lookahead, &mut ChangeSet::default());
         }
     }
 }

crates/chain/tests/test_indexed_tx_graph.rs

@@ -219,7 +219,7 @@ fn insert_relevant_txs() {
     let (descriptor, _) = Descriptor::parse_descriptor(&Secp256k1::signing_only(), DESCRIPTORS[0])
         .expect("must be valid");
     let spk_0 = descriptor.at_derivation_index(0).unwrap().script_pubkey();
-    let spk_1 = descriptor.at_derivation_index(9).unwrap().script_pubkey();
+    let spk_9 = descriptor.at_derivation_index(9).unwrap().script_pubkey();
     let lookahead = 10;
 
     let mut graph = IndexedTxGraph::<ConfirmationBlockTime, KeychainTxOutIndex<()>>::new({
@@ -237,7 +237,7 @@ fn insert_relevant_txs() {
             },
             TxOut {
                 value: Amount::from_sat(20_000),
-                script_pubkey: spk_1,
+                script_pubkey: spk_9,
             },
         ],
         ..new_tx(0)
@@ -269,12 +269,12 @@ fn insert_relevant_txs() {
         indexer: keychain_txout::ChangeSet {
             last_revealed: [(descriptor.descriptor_id(), 9_u32)].into(),
             spk_cache: [(descriptor.descriptor_id(), {
-                let index_after_spk_1 = 9 /* index of spk_1 */ + 1;
+                let index_after_spk_9 = 9 /* index of spk_9 */ + 1;
                 SpkIterator::new_with_range(
                     &descriptor,
                     // This will also persist the staged spk cache inclusions from prev call to
                     // `.insert_descriptor`.
-                    0..index_after_spk_1 + lookahead,
+                    index_after_spk_9..index_after_spk_9 + lookahead,
                 )
                 .collect()
             })]
@@ -294,7 +294,7 @@ fn insert_relevant_txs() {
             last_revealed: changeset.indexer.last_revealed,
             spk_cache: [(
                 descriptor.descriptor_id(),
-                SpkIterator::new_with_range(&descriptor, 0..=9 /* index of spk_1*/  + lookahead)
+                SpkIterator::new_with_range(&descriptor, 0..=9 /* index of spk_9*/  + lookahead)
                     .collect(),
             )]
             .into(),