feat(chain)!: Introduce descriptor in keychain...

danielabrozzoni · danielabrozzoni · commit c14504e97532 · 2023-11-13T10:03:13.000+01:00
...txout index changeset Fixes #1101
diff --git a/crates/chain/src/keychain/txout_index.rs b/crates/chain/src/keychain/txout_index.rs
@@ -12,7 +12,7 @@ use core::{fmt::Debug, ops::Deref};
 use crate::Append;
 
 /// Represents updates to the derivation index of a [`KeychainTxOutIndex`].
-/// It maps each keychain `K` to its last revealed index.
+/// It maps each keychain `K` to a descriptor and its last revealed index.
 ///
 /// It can be applied to [`KeychainTxOutIndex`] with [`apply_changeset`]. [`ChangeSet] are
 /// monotone in that they will never decrease the revealed derivation index.
@@ -32,26 +32,32 @@ use crate::Append;
     )
 )]
 #[must_use]
-pub struct ChangeSet<K>(pub BTreeMap<K, u32>);
+pub struct ChangeSet<K>(pub BTreeMap<K, (Descriptor<DescriptorPublicKey>, u32)>);
 
 impl<K> ChangeSet<K> {
     /// Get the inner map of the keychain to its new derivation index.
-    pub fn as_inner(&self) -> &BTreeMap<K, u32> {
+    pub fn as_inner(&self) -> &BTreeMap<K, (Descriptor<DescriptorPublicKey>, u32)> {
         &self.0
     }
 }
 
 impl<K: Ord> Append for ChangeSet<K> {
     /// Append another [`ChangeSet`] into self.
     ///
-    /// If the keychain already exists, increase the index when the other's index > self's index.
-    /// If the keychain did not exist, append the new keychain.
+    /// For each keychain in the given [`ChangeSet`]:
+    /// - If the keychain already exists in the old [`ChangeSet`], and the two descriptors are the
+    /// same, we increase the index when new index > old index.
+    /// - If the keychain already exists, but with a different descriptor, we overwrite the
+    /// original descriptor with the changeset one.
+    /// - If the keychain did not exist, append the new keychain.
     fn append(&mut self, mut other: Self) {
-        self.0.iter_mut().for_each(|(key, index)| {
-            if let Some(other_index) = other.0.remove(key) {
-                *index = other_index.max(*index);
+        for (key, (descriptor, index)) in &mut self.0 {
+            if let Some((other_descriptor, other_index)) = other.0.remove(key) {
+                if other_descriptor == *descriptor {
+                    *index = other_index.max(*index);
+                }
             }
-        });
+        }
 
         self.0.append(&mut other.0);
     }
@@ -68,8 +74,8 @@ impl<K> Default for ChangeSet<K> {
     }
 }
 
-impl<K> AsRef<BTreeMap<K, u32>> for ChangeSet<K> {
-    fn as_ref(&self) -> &BTreeMap<K, u32> {
+impl<K> AsRef<BTreeMap<K, (Descriptor<DescriptorPublicKey>, u32)>> for ChangeSet<K> {
+    fn as_ref(&self) -> &BTreeMap<K, (Descriptor<DescriptorPublicKey>, u32)> {
         &self.0
     }
 }
@@ -169,7 +175,16 @@ impl<K: Clone + Ord + Debug> Indexer for KeychainTxOutIndex<K> {
     }
 
     fn initial_changeset(&self) -> Self::ChangeSet {
-        super::ChangeSet(self.last_revealed.clone())
+        let changeset = self
+            .keychains
+            .clone()
+            .into_iter()
+            .map(|(key, descriptor)| {
+                let index = self.last_revealed.get(&key).expect("Must be here");
+                (key, (descriptor, *index))
+            })
+            .collect();
+        super::ChangeSet(changeset)
     }
 
     fn apply_changeset(&mut self, changeset: Self::ChangeSet) {
@@ -488,7 +503,9 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
                 debug_assert!(_old_index < Some(index));
                 (
                     SpkIterator::new_with_range(descriptor.clone(), next_reveal_index..index + 1),
-                    super::ChangeSet(core::iter::once((keychain.clone(), index)).collect()),
+                    super::ChangeSet(
+                        core::iter::once((keychain.clone(), (descriptor.clone(), index))).collect(),
+                    ),
                 )
             }
             None => (
@@ -621,6 +638,12 @@ impl<K: Clone + Ord + Debug> KeychainTxOutIndex<K> {
     /// Applies the derivation changeset to the [`KeychainTxOutIndex`], extending the number of
     /// derived scripts per keychain, as specified in the `changeset`.
     pub fn apply_changeset(&mut self, changeset: super::ChangeSet<K>) {
-        let _ = self.reveal_to_target_multi(&changeset.0);
+        let _ = self.reveal_to_target_multi(
+            &changeset
+                .0
+                .into_iter()
+                .map(|(key, (_, index))| (key, index))
+                .collect(),
+        );
     }
 }
diff --git a/crates/chain/tests/test_indexed_tx_graph.rs b/crates/chain/tests/test_indexed_tx_graph.rs
@@ -30,7 +30,7 @@ fn insert_relevant_txs() {
     let spk_1 = descriptor.at_derivation_index(9).unwrap().script_pubkey();
 
     let mut graph = IndexedTxGraph::<ConfirmationHeightAnchor, KeychainTxOutIndex<()>>::default();
-    graph.index.add_keychain((), descriptor);
+    graph.index.add_keychain((), descriptor.clone());
     graph.index.set_lookahead(&(), 10);
 
     let tx_a = Transaction {
@@ -70,7 +70,7 @@ fn insert_relevant_txs() {
             txs: txs.clone().into(),
             ..Default::default()
         },
-        indexer: keychain::ChangeSet([((), 9_u32)].into()),
+        indexer: keychain::ChangeSet([((), (descriptor, 9_u32))].into()),
     };
 
     assert_eq!(
diff --git a/crates/chain/tests/test_keychain_txout_index.rs b/crates/chain/tests/test_keychain_txout_index.rs
@@ -46,52 +46,91 @@ fn spk_at_index(descriptor: &Descriptor<DescriptorPublicKey>, index: u32) -> Scr
 fn append_keychain_derivation_indices() {
     #[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Debug)]
     enum Keychain {
+        Zero,
         One,
         Two,
         Three,
         Four,
     }
-    let mut lhs_di = BTreeMap::<Keychain, u32>::default();
-    let mut rhs_di = BTreeMap::<Keychain, u32>::default();
-    lhs_di.insert(Keychain::One, 7);
-    lhs_di.insert(Keychain::Two, 0);
-    rhs_di.insert(Keychain::One, 3);
-    rhs_di.insert(Keychain::Two, 5);
-    lhs_di.insert(Keychain::Three, 3);
-    rhs_di.insert(Keychain::Four, 4);
+
+    let descriptors: Vec<_> = common::DESCRIPTORS
+        .into_iter()
+        .map(|s| {
+            Descriptor::parse_descriptor(&Secp256k1::signing_only(), s)
+                .unwrap()
+                .0
+        })
+        .collect();
+
+    let mut lhs_di = BTreeMap::<Keychain, (Descriptor<DescriptorPublicKey>, u32)>::default();
+    let mut rhs_di = BTreeMap::<Keychain, (Descriptor<DescriptorPublicKey>, u32)>::default();
+    lhs_di.insert(Keychain::Zero, (descriptors[0].clone(), 1));
+    lhs_di.insert(Keychain::One, (descriptors[1].clone(), 7));
+    lhs_di.insert(Keychain::Two, (descriptors[2].clone(), 0));
+
+    lhs_di.insert(Keychain::Zero, (descriptors[1].clone(), 13));
+    rhs_di.insert(Keychain::One, (descriptors[1].clone(), 3));
+    rhs_di.insert(Keychain::Two, (descriptors[2].clone(), 5));
+    lhs_di.insert(Keychain::Three, (descriptors[3].clone(), 3));
+    rhs_di.insert(Keychain::Four, (descriptors[4].clone(), 4));
 
     let mut lhs = ChangeSet(lhs_di);
     let rhs = ChangeSet(rhs_di);
     lhs.append(rhs);
 
+    // Descriptor gets updated
+    assert_eq!(
+        lhs.0.get(&Keychain::Zero),
+        (Some(&(descriptors[1].clone(), 13)))
+    );
     // Exiting index doesn't update if the new index in `other` is lower than `self`.
-    assert_eq!(lhs.0.get(&Keychain::One), Some(&7));
+    assert_eq!(
+        lhs.0.get(&Keychain::One),
+        (Some(&(descriptors[1].clone(), 7)))
+    );
     // Existing index updates if the new index in `other` is higher than `self`.
-    assert_eq!(lhs.0.get(&Keychain::Two), Some(&5));
+    assert_eq!(
+        lhs.0.get(&Keychain::Two),
+        Some(&(descriptors[2].clone(), 5))
+    );
     // Existing index is unchanged if keychain doesn't exist in `other`.
-    assert_eq!(lhs.0.get(&Keychain::Three), Some(&3));
+    assert_eq!(
+        lhs.0.get(&Keychain::Three),
+        Some(&(descriptors[3].clone(), 3))
+    );
     // New keychain gets added if the keychain is in `other` but not in `self`.
-    assert_eq!(lhs.0.get(&Keychain::Four), Some(&4));
+    assert_eq!(
+        lhs.0.get(&Keychain::Four),
+        Some(&(descriptors[4].clone(), 4))
+    );
 }
 
 #[test]
 fn test_set_all_derivation_indices() {
     use bdk_chain::indexed_tx_graph::Indexer;
 
-    let (mut txout_index, _, _) = init_txout_index();
-    let derive_to: BTreeMap<_, _> =
-        [(TestKeychain::External, 12), (TestKeychain::Internal, 24)].into();
+    let (mut txout_index, external_desc, internal_desc) = init_txout_index();
+    let result: BTreeMap<_, _> = [
+        (TestKeychain::External, (external_desc, 12)),
+        (TestKeychain::Internal, (internal_desc, 24)),
+    ]
+    .into();
+    let derive_to = result
+        .clone()
+        .into_iter()
+        .map(|(k, (_, i))| (k, i))
+        .collect();
     assert_eq!(
         txout_index.reveal_to_target_multi(&derive_to).1.as_inner(),
-        &derive_to
+        &result
     );
     assert_eq!(txout_index.last_revealed_indices(), &derive_to);
     assert_eq!(
         txout_index.reveal_to_target_multi(&derive_to).1,
         keychain::ChangeSet::default(),
         "no changes if we set to the same thing"
     );
-    assert_eq!(txout_index.initial_changeset().as_inner(), &derive_to);
+    assert_eq!(txout_index.initial_changeset().as_inner(), &result);
 }
 
 #[test]
@@ -123,7 +162,7 @@ fn test_lookahead() {
         );
         assert_eq!(
             revealed_changeset.as_inner(),
-            &[(TestKeychain::External, index)].into()
+            &[(TestKeychain::External, (external_desc.clone(), index))].into()
         );
 
         assert_eq!(
@@ -175,7 +214,7 @@ fn test_lookahead() {
     );
     assert_eq!(
         revealed_changeset.as_inner(),
-        &[(TestKeychain::Internal, 24)].into()
+        &[(TestKeychain::Internal, (internal_desc.clone(), 24))].into()
     );
     assert_eq!(
         txout_index.inner().all_spks().len(),
@@ -284,7 +323,7 @@ fn test_scan_with_lookahead() {
         let changeset = txout_index.index_txout(op, &txout);
         assert_eq!(
             changeset.as_inner(),
-            &[(TestKeychain::External, spk_i)].into()
+            &[(TestKeychain::External, (external_desc.clone(), spk_i))].into()
         );
         assert_eq!(
             txout_index.last_revealed_index(&TestKeychain::External),
@@ -328,7 +367,7 @@ fn test_wildcard_derivations() {
     assert_eq!(txout_index.next_index(&TestKeychain::External), (0, true));
     let (spk, changeset) = txout_index.reveal_next_spk(&TestKeychain::External);
     assert_eq!(spk, (0_u32, external_spk_0.as_script()));
-    assert_eq!(changeset.as_inner(), &[(TestKeychain::External, 0)].into());
+    assert_eq!(changeset.as_inner(), &[(TestKeychain::External, (external_desc.clone(), 0))].into());
     let (spk, changeset) = txout_index.next_unused_spk(&TestKeychain::External);
     assert_eq!(spk, (0_u32, external_spk_0.as_script()));
     assert_eq!(changeset.as_inner(), &[].into());
@@ -352,7 +391,7 @@ fn test_wildcard_derivations() {
     let (spk, changeset) = txout_index.reveal_next_spk(&TestKeychain::External);
     assert_eq!(spk, (26, external_spk_26.as_script()));
 
-    assert_eq!(changeset.as_inner(), &[(TestKeychain::External, 26)].into());
+    assert_eq!(changeset.as_inner(), &[(TestKeychain::External, (external_desc.clone(), 26))].into());
 
     let (spk, changeset) = txout_index.next_unused_spk(&TestKeychain::External);
     assert_eq!(spk, (16, external_spk_16.as_script()));
@@ -366,7 +405,7 @@ fn test_wildcard_derivations() {
 
     let (spk, changeset) = txout_index.next_unused_spk(&TestKeychain::External);
     assert_eq!(spk, (27, external_spk_27.as_script()));
-    assert_eq!(changeset.as_inner(), &[(TestKeychain::External, 27)].into());
+    assert_eq!(changeset.as_inner(), &[(TestKeychain::External, (external_desc, 27))].into());
 }
 
 #[test]
@@ -380,7 +419,7 @@ fn test_non_wildcard_derivations() {
         .unwrap()
         .script_pubkey();
 
-    txout_index.add_keychain(TestKeychain::External, no_wildcard_descriptor);
+    txout_index.add_keychain(TestKeychain::External, no_wildcard_descriptor.clone());
 
     // given:
     // - `txout_index` with no stored scripts
@@ -391,7 +430,10 @@ fn test_non_wildcard_derivations() {
     assert_eq!(txout_index.next_index(&TestKeychain::External), (0, true));
     let (spk, changeset) = txout_index.reveal_next_spk(&TestKeychain::External);
     assert_eq!(spk, (0, external_spk.as_script()));
-    assert_eq!(changeset.as_inner(), &[(TestKeychain::External, 0)].into());
+    assert_eq!(
+        changeset.as_inner(),
+        &[(TestKeychain::External, (no_wildcard_descriptor, 0))].into()
+    );
 
     let (spk, changeset) = txout_index.next_unused_spk(&TestKeychain::External);
     assert_eq!(spk, (0, external_spk.as_script()));
