Skip to content

Commit

Permalink
feat(chain)!: Introduce descriptor in keychain...
Browse files Browse the repository at this point in the history 
...txout index changeset

Fixes #1101
  • Loading branch information
@danielabrozzoni
danielabrozzoni committed Nov 10, 2023
1 parent a97d662 commit 7c4fea5
Show file tree
Hide file tree
Showing 3 changed files with 104 additions and 39 deletions.
46 changes: 34 additions & 12 deletions crates/chain/src/keychain/txout_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,26 +32,31 @@ 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 already exists, and the two descriptors are the same, increase the index
/// when the other's index > self's 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);
}
Expand All @@ -68,8 +73,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
}
}
Expand Down Expand Up @@ -169,7 +174,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) {
Expand Down Expand Up @@ -488,7 +502,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 => (
Expand Down Expand Up @@ -621,6 +637,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(),
);
}
}
4 changes: 2 additions & 2 deletions crates/chain/tests/test_indexed_tx_graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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!(
Expand Down
93 changes: 68 additions & 25 deletions crates/chain/tests/test_keychain_txout_index.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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!(
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -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),
Expand All @@ -298,6 +337,7 @@ fn test_scan_with_lookahead() {

// now try with index 41 (lookahead surpassed), we expect that the txout to not be indexed
let spk_41 = external_desc
.clone()
.at_derivation_index(41)
.unwrap()
.script_pubkey();
Expand Down Expand Up @@ -328,7 +368,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());
Expand All @@ -352,7 +392,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()));
Expand All @@ -366,7 +406,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]
Expand All @@ -380,7 +420,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
Expand All @@ -391,7 +431,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()));
Expand Down

0 comments on commit 7c4fea5

Please sign in to comment.