Skip to content

Commit

Permalink
feat(chain): get rid of TxGraph::determine_changeset
Browse files Browse the repository at this point in the history 
Contain most of the insertion logic in `.insert_{}` methods, thus
simplifying `.apply_{}` methods. We can also get rid of
`.determine_changeset`.
  • Loading branch information
@evanlinjin
evanlinjin committed Aug 22, 2024
1 parent 6008897 commit 71a3e0e
Show file tree
Hide file tree
Showing 2 changed files with 112 additions and 167 deletions.
224 changes: 100 additions & 124 deletions crates/chain/src/tx_graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -512,28 +512,66 @@ impl<A: Clone + Ord> TxGraph<A> {
///
/// [`apply_changeset`]: Self::apply_changeset
pub fn insert_txout(&mut self, outpoint: OutPoint, txout: TxOut) -> ChangeSet<A> {
let mut update = Self::default();
update.txs.insert(
outpoint.txid,
(
TxNodeInternal::Partial([(outpoint.vout, txout)].into()),
BTreeSet::new(),
),
);
self.apply_update(update)
let mut changeset = ChangeSet::<A>::default();
let (tx_node, _) = self.txs.entry(outpoint.txid).or_default();
match tx_node {
TxNodeInternal::Whole(_) => {
// ignore this txout we have the full one already.
// NOTE: You might think putting a debug_assert! here to check the output being
// replaced was actually correct is a good idea but the tests have already been
// written assuming this never panics.
}
TxNodeInternal::Partial(partial_tx) => {
match partial_tx.insert(outpoint.vout, txout.clone()) {
Some(old_txout) => {
debug_assert_eq!(
txout, old_txout,
"txout of the same outpoint should never change"
);
}
None => {
changeset.txouts.insert(outpoint, txout);
}
}
}
}
changeset
}

/// Inserts the given transaction into [`TxGraph`].
///
/// The [`ChangeSet`] returned will be empty if `tx` already exists.
pub fn insert_tx<T: Into<Arc<Transaction>>>(&mut self, tx: T) -> ChangeSet<A> {
let tx = tx.into();
let mut update = Self::default();
update.txs.insert(
tx.compute_txid(),
(TxNodeInternal::Whole(tx), BTreeSet::new()),
);
self.apply_update(update)
let tx: Arc<Transaction> = tx.into();
let txid = tx.compute_txid();
let mut changeset = ChangeSet::<A>::default();

let (tx_node, _) = self.txs.entry(txid).or_default();
match tx_node {
TxNodeInternal::Whole(existing_tx) => {
debug_assert_eq!(
existing_tx.as_ref(),
tx.as_ref(),
"tx of same txid should never change"
);
}
partial_tx => {
for txin in &tx.input {
// this means the tx is coinbase so there is no previous output
if txin.previous_output.is_null() {
continue;
}
self.spends
.entry(txin.previous_output)
.or_default()
.insert(txid);
}
*partial_tx = TxNodeInternal::Whole(tx.clone());
changeset.txs.insert(tx);
}
}

changeset
}

/// Batch insert unconfirmed transactions.
Expand All @@ -558,9 +596,17 @@ impl<A: Clone + Ord> TxGraph<A> {
/// The [`ChangeSet`] returned will be empty if graph already knows that `txid` exists in
/// `anchor`.
pub fn insert_anchor(&mut self, txid: Txid, anchor: A) -> ChangeSet<A> {
let mut update = Self::default();
update.anchors.insert((anchor, txid));
self.apply_update(update)
let mut changeset = ChangeSet::<A>::default();
if self.anchors.insert((anchor.clone(), txid)) {
let (_tx_node, anchors) = self.txs.entry(txid).or_default();
let _inserted = anchors.insert(anchor.clone());
debug_assert!(
_inserted,
"anchors in `.anchors` and `.txs` should be consistent"
);
changeset.anchors.insert((anchor, txid));
}
changeset
}

/// Inserts the given `seen_at` for `txid` into [`TxGraph`].
Expand All @@ -571,9 +617,13 @@ impl<A: Clone + Ord> TxGraph<A> {
///
/// [`update_last_seen_unconfirmed`]: Self::update_last_seen_unconfirmed
pub fn insert_seen_at(&mut self, txid: Txid, seen_at: u64) -> ChangeSet<A> {
let mut update = Self::default();
update.last_seen.insert(txid, seen_at);
self.apply_update(update)
let mut changeset = ChangeSet::<A>::default();
let last_seen = self.last_seen.entry(txid).or_default();
if seen_at > *last_seen {
*last_seen = seen_at;
changeset.last_seen.insert(txid, seen_at);
}
changeset
}

/// Update the last seen time for all unconfirmed transactions.
Expand Down Expand Up @@ -641,124 +691,50 @@ impl<A: Clone + Ord> TxGraph<A> {
/// The returned [`ChangeSet`] is the set difference between `update` and `self` (transactions that
/// exist in `update` but not in `self`).
pub fn apply_update(&mut self, update: TxGraph<A>) -> ChangeSet<A> {
let changeset = self.determine_changeset(update);
self.apply_changeset(changeset.clone());
let mut changeset = ChangeSet::<A>::default();
for tx_node in update.full_txs() {
changeset.merge(self.insert_tx(tx_node.tx));
}
for (outpoint, txout) in update.floating_txouts() {
changeset.merge(self.insert_txout(outpoint, txout.clone()));
}
for (anchor, txid) in &update.anchors {
changeset.merge(self.insert_anchor(*txid, anchor.clone()));
}
for (&txid, &last_seen) in &update.last_seen {
changeset.merge(self.insert_seen_at(txid, last_seen));
}
changeset
}

/// Determines the [`ChangeSet`] between `self` and an empty [`TxGraph`].
pub fn initial_changeset(&self) -> ChangeSet<A> {
Self::default().determine_changeset(self.clone())
ChangeSet {
txs: self.full_txs().map(|tx_node| tx_node.tx).collect(),
txouts: self
.floating_txouts()
.map(|(op, txout)| (op, txout.clone()))
.collect(),
anchors: self.anchors.clone(),
last_seen: self.last_seen.iter().map(|(&k, &v)| (k, v)).collect(),
}
}

/// Applies [`ChangeSet`] to [`TxGraph`].
pub fn apply_changeset(&mut self, changeset: ChangeSet<A>) {
for wrapped_tx in changeset.txs {
let tx = wrapped_tx.as_ref();
let txid = tx.compute_txid();

tx.input
.iter()
.map(|txin| txin.previous_output)
// coinbase spends are not to be counted
.filter(|outpoint| !outpoint.is_null())
// record spend as this tx has spent this outpoint
.for_each(|outpoint| {
self.spends.entry(outpoint).or_default().insert(txid);
});

match self.txs.get_mut(&txid) {
Some((tx_node @ TxNodeInternal::Partial(_), _)) => {
*tx_node = TxNodeInternal::Whole(wrapped_tx.clone());
}
Some((TxNodeInternal::Whole(tx), _)) => {
debug_assert_eq!(
tx.as_ref().compute_txid(),
txid,
"tx should produce txid that is same as key"
);
}
None => {
self.txs
.insert(txid, (TxNodeInternal::Whole(wrapped_tx), BTreeSet::new()));
}
}
for tx in changeset.txs {
let _ = self.insert_tx(tx);
}

for (outpoint, txout) in changeset.txouts {
let tx_entry = self.txs.entry(outpoint.txid).or_default();

match tx_entry {
(TxNodeInternal::Whole(_), _) => { /* do nothing since we already have full tx */ }
(TxNodeInternal::Partial(txouts), _) => {
txouts.insert(outpoint.vout, txout);
}
}
let _ = self.insert_txout(outpoint, txout);
}

for (anchor, txid) in changeset.anchors {
if self.anchors.insert((anchor.clone(), txid)) {
let (_, anchors) = self.txs.entry(txid).or_default();
anchors.insert(anchor);
}
let _ = self.insert_anchor(txid, anchor);
}

for (txid, new_last_seen) in changeset.last_seen {
let last_seen = self.last_seen.entry(txid).or_default();
if new_last_seen > *last_seen {
*last_seen = new_last_seen;
}
for (txid, seen_at) in changeset.last_seen {
let _ = self.insert_seen_at(txid, seen_at);
}
}

/// Previews the resultant [`ChangeSet`] when [`Self`] is updated against the `update` graph.
///
/// The [`ChangeSet`] would be the set difference between `update` and `self` (transactions that
/// exist in `update` but not in `self`).
pub(crate) fn determine_changeset(&self, update: TxGraph<A>) -> ChangeSet<A> {
let mut changeset = ChangeSet::<A>::default();

for (&txid, (update_tx_node, _)) in &update.txs {
match (self.txs.get(&txid), update_tx_node) {
(None, TxNodeInternal::Whole(update_tx)) => {
changeset.txs.insert(update_tx.clone());
}
(None, TxNodeInternal::Partial(update_txos)) => {
changeset.txouts.extend(
update_txos
.iter()
.map(|(&vout, txo)| (OutPoint::new(txid, vout), txo.clone())),
);
}
(Some((TxNodeInternal::Whole(_), _)), _) => {}
(Some((TxNodeInternal::Partial(_), _)), TxNodeInternal::Whole(update_tx)) => {
changeset.txs.insert(update_tx.clone());
}
(
Some((TxNodeInternal::Partial(txos), _)),
TxNodeInternal::Partial(update_txos),
) => {
changeset.txouts.extend(
update_txos
.iter()
.filter(|(vout, _)| !txos.contains_key(*vout))
.map(|(&vout, txo)| (OutPoint::new(txid, vout), txo.clone())),
);
}
}
}

for (txid, update_last_seen) in update.last_seen {
let prev_last_seen = self.last_seen.get(&txid).copied();
if Some(update_last_seen) > prev_last_seen {
changeset.last_seen.insert(txid, update_last_seen);
}
}

changeset.anchors = update.anchors.difference(&self.anchors).cloned().collect();

changeset
}
}

impl<A: Anchor> TxGraph<A> {
Expand Down
55 changes: 12 additions & 43 deletions crates/chain/tests/test_tx_graph.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -301,59 +301,28 @@ fn insert_tx_can_retrieve_full_tx_from_graph() {
#[test]
fn insert_tx_displaces_txouts() {
let mut tx_graph = TxGraph::<()>::default();

let tx = Transaction {
version: transaction::Version::ONE,
lock_time: absolute::LockTime::ZERO,
input: vec![],
output: vec![TxOut {
value: Amount::from_sat(42_000),
script_pubkey: ScriptBuf::new(),
script_pubkey: ScriptBuf::default(),
}],
};
let txid = tx.compute_txid();
let outpoint = OutPoint::new(txid, 0);
let txout = tx.output.first().unwrap();

let changeset = tx_graph.insert_txout(
OutPoint {
txid: tx.compute_txid(),
vout: 0,
},
TxOut {
value: Amount::from_sat(1_337_000),
script_pubkey: ScriptBuf::default(),
},
);

let changeset = tx_graph.insert_txout(outpoint, txout.clone());
assert!(!changeset.is_empty());

let _ = tx_graph.insert_txout(
OutPoint {
txid: tx.compute_txid(),
vout: 0,
},
TxOut {
value: Amount::from_sat(1_000_000_000),
script_pubkey: ScriptBuf::new(),
},
);

let _changeset = tx_graph.insert_tx(tx.clone());

assert_eq!(
tx_graph
.get_txout(OutPoint {
txid: tx.compute_txid(),
vout: 0
})
.unwrap()
.value,
Amount::from_sat(42_000)
);
assert_eq!(
tx_graph.get_txout(OutPoint {
txid: tx.compute_txid(),
vout: 1
}),
None
);
let changeset = tx_graph.insert_tx(tx.clone());
assert_eq!(changeset.txs.len(), 1);
assert!(changeset.txouts.is_empty());
assert!(tx_graph.get_tx(txid).is_some());
assert_eq!(tx_graph.get_txout(outpoint), Some(txout));
}

#[test]
Expand Down Expand Up @@ -385,7 +354,7 @@ fn insert_txout_does_not_displace_tx() {
let _ = tx_graph.insert_txout(
OutPoint {
txid: tx.compute_txid(),
vout: 0,
vout: 1,
},
TxOut {
value: Amount::from_sat(1_000_000_000),
Expand Down

0 comments on commit 71a3e0e

Please sign in to comment.