Skip to content

Commit

Permalink
feat(chain): add map_anchors for TxGraph and ChangeSet
Browse files Browse the repository at this point in the history
  • Loading branch information
yanganto committed Feb 13, 2024
1 parent 022d5a2 commit 5489f90
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 2 deletions.
35 changes: 35 additions & 0 deletions crates/chain/src/tx_graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -451,6 +451,21 @@ impl<A> TxGraph<A> {
}
}

impl<A: Clone + Ord> TxGraph<A> {
/// Transform the [`TxGraph`] to have [`Anchor`]s of another type.
///
/// This takes in a closure of signature `FnMut(A) -> A2` which is called for each [`Anchor`] to
/// transform it.
pub fn map_anchors<A2: Clone + Ord, F>(self, f: F) -> TxGraph<A2>
where
F: FnMut(A) -> A2,
{
let mut new_graph = TxGraph::<A2>::default();
new_graph.apply_changeset(self.initial_changeset().map_anchors(f));
new_graph
}
}

impl<A: Clone + Ord> TxGraph<A> {
/// Construct a new [`TxGraph`] from a list of transactions.
pub fn new(txs: impl IntoIterator<Item = Transaction>) -> Self {
Expand Down Expand Up @@ -1296,6 +1311,26 @@ impl<A: Ord> Append for ChangeSet<A> {
}
}

impl<A: Ord> ChangeSet<A> {
/// Transform the [`ChangeSet`] to have [`Anchor`]s of another type.
///
/// This takes in a closure of signature `FnMut(A) -> A2` which is called for each [`Anchor`] to
/// transform it.
pub fn map_anchors<A2: Ord, F>(self, mut f: F) -> ChangeSet<A2>
where
F: FnMut(A) -> A2,
{
ChangeSet {
txs: self.txs,
txouts: self.txouts,
anchors: BTreeSet::<(A2, Txid)>::from_iter(
self.anchors.into_iter().map(|(a, txid)| (f(a), txid)),
),
last_seen: self.last_seen,
}
}
}

impl<A> AsRef<TxGraph<A>> for TxGraph<A> {
fn as_ref(&self) -> &TxGraph<A> {
self
Expand Down
4 changes: 2 additions & 2 deletions crates/chain/tests/common/tx_template.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ impl TxOutTemplate {
}

#[allow(dead_code)]
pub fn init_graph<'a, A: Anchor + Copy + 'a>(
pub fn init_graph<'a, A: Anchor + Clone + 'a>(
tx_templates: impl IntoIterator<Item = &'a TxTemplate<'a, A>>,
) -> (TxGraph<A>, SpkTxOutIndex<u32>, HashMap<&'a str, Txid>) {
let (descriptor, _) = Descriptor::parse_descriptor(&Secp256k1::signing_only(), "tr(tprv8ZgxMBicQKsPd3krDUsBAmtnRsK3rb8u5yi1zhQgMhF1tR8MW7xfE4rnrbbsrbPR52e7rKapu6ztw1jXveJSCGHEriUGZV7mCe88duLp5pj/86'/1'/0'/0/*)").unwrap();
Expand Down Expand Up @@ -126,7 +126,7 @@ pub fn init_graph<'a, A: Anchor + Copy + 'a>(
spk_index.scan(&tx);
let _ = graph.insert_tx(tx.clone());
for anchor in tx_tmp.anchors.iter() {
let _ = graph.insert_anchor(tx.txid(), *anchor);
let _ = graph.insert_anchor(tx.txid(), anchor.clone());
}
if let Some(seen_at) = tx_tmp.last_seen {
let _ = graph.insert_seen_at(tx.txid(), seen_at);
Expand Down
85 changes: 85 additions & 0 deletions crates/chain/tests/test_tx_graph.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,9 @@ use bdk_chain::{
use bitcoin::{
absolute, hashes::Hash, BlockHash, OutPoint, ScriptBuf, Transaction, TxIn, TxOut, Txid,
};
use common::*;
use core::iter;
use rand::RngCore;
use std::vec;

#[test]
Expand Down Expand Up @@ -1172,3 +1174,86 @@ fn test_missing_blocks() {
),
]);
}

#[test]
/// The `map_anchors` allow a caller to pass a function to reconstruct the [`TxGraph`] with any [`Anchor`],
/// even though the function is non-deterministic.
fn call_map_anchors_with_non_deterministic_anchor() {
#[derive(Debug, Default, Clone, PartialEq, Eq, Copy, PartialOrd, Ord, core::hash::Hash)]
/// A non-deterministic anchor
pub struct NonDeterministicAnchor {
pub anchor_block: BlockId,
pub non_deterministic_field: u32,
}

let template = [
TxTemplate {
tx_name: "tx1",
inputs: &[TxInTemplate::Bogus],
outputs: &[TxOutTemplate::new(10000, Some(1))],
anchors: &[block_id!(1, "A")],
last_seen: None,
},
TxTemplate {
tx_name: "tx2",
inputs: &[TxInTemplate::PrevTx("tx1", 0)],
outputs: &[TxOutTemplate::new(20000, Some(2))],
anchors: &[block_id!(2, "B")],
..Default::default()
},
TxTemplate {
tx_name: "tx3",
inputs: &[TxInTemplate::PrevTx("tx2", 0)],
outputs: &[TxOutTemplate::new(30000, Some(3))],
anchors: &[block_id!(3, "C"), block_id!(4, "D")],
..Default::default()
},
];
let (graph, _, _) = init_graph(&template);
let new_graph = graph.clone().map_anchors(|a| NonDeterministicAnchor {
anchor_block: a,
// A non-deterministic value
non_deterministic_field: rand::thread_rng().next_u32(),
});

// Check all the details in new_graph reconstruct as well

let mut full_txs_vec: Vec<_> = graph.full_txs().collect();
full_txs_vec.sort();
let mut new_txs_vec: Vec<_> = new_graph.full_txs().collect();
new_txs_vec.sort();
let mut new_txs = new_txs_vec.iter();

for tx_node in full_txs_vec.iter() {
let new_txnode = new_txs.next().unwrap();
assert_eq!(new_txnode.txid, tx_node.txid);
assert_eq!(new_txnode.tx, tx_node.tx);
assert_eq!(
new_txnode.last_seen_unconfirmed,
tx_node.last_seen_unconfirmed
);
assert_eq!(new_txnode.anchors.len(), tx_node.anchors.len());

let mut new_anchors: Vec<_> = new_txnode.anchors.iter().map(|a| a.anchor_block).collect();
new_anchors.sort();
let mut old_anchors: Vec<_> = tx_node.anchors.iter().copied().collect();
old_anchors.sort();
assert_eq!(new_anchors, old_anchors);
}
assert!(new_txs.next().is_none());

let new_graph_anchors: Vec<_> = new_graph
.all_anchors()
.iter()
.map(|i| i.0.anchor_block)
.collect();
assert_eq!(
new_graph_anchors,
vec![
block_id!(1, "A"),
block_id!(2, "B"),
block_id!(3, "C"),
block_id!(4, "D"),
]
);
}

0 comments on commit 5489f90

Please sign in to comment.