[bdk_chain_redesign] Fix tx_graph::Additions::append logic

evanlinjin · evanlinjin · commit e536307e5c55 · 2023-04-28T18:54:36.000+08:00
* `Additions` now implements `Append` and uses `Append` to implement
  `append()`.
* `append()` logic enforces that `last_seen` values should only
  increase.
* Test written for `append()` with `last_seen` behaviour.
diff --git a/crates/chain/src/chain_graph.rs b/crates/chain/src/chain_graph.rs
@@ -3,7 +3,7 @@ use crate::{
     collections::HashSet,
     sparse_chain::{self, ChainPosition, SparseChain},
     tx_graph::{self, TxGraph},
-    BlockId, ForEachTxOut, FullTxOut, TxHeight,
+    Append, BlockId, ForEachTxOut, FullTxOut, TxHeight,
 };
 use alloc::{string::ToString, vec::Vec};
 use bitcoin::{OutPoint, Transaction, TxOut, Txid};
diff --git a/crates/chain/src/tx_graph.rs b/crates/chain/src/tx_graph.rs
@@ -55,7 +55,9 @@
 //! assert!(additions.is_empty());
 //! ```
 
-use crate::{collections::*, Anchor, BlockId, ChainOracle, ForEachTxOut, FullTxOut, ObservedAs};
+use crate::{
+    collections::*, Anchor, Append, BlockId, ChainOracle, ForEachTxOut, FullTxOut, ObservedAs,
+};
 use alloc::vec::Vec;
 use bitcoin::{OutPoint, Transaction, TxOut, Txid};
 use core::{
@@ -112,17 +114,6 @@ impl<'a, T, A> Deref for TxNode<'a, T, A> {
     }
 }
 
-impl<'a, A> TxNode<'a, Transaction, A> {
-    pub fn from_tx(tx: &'a Transaction, anchors: &'a BTreeSet<A>) -> Self {
-        Self {
-            txid: tx.txid(),
-            tx,
-            anchors,
-            last_seen_unconfirmed: 0,
-        }
-    }
-}
-
 /// Internal representation of a transaction node of a [`TxGraph`].
 ///
 /// This can either be a whole transaction, or a partial transaction (where we only have select
@@ -602,7 +593,7 @@ impl<A: Clone + Ord> TxGraph<A> {
 
 impl<A: Anchor> TxGraph<A> {
     /// Get all heights that are relevant to the graph.
-    pub fn relevant_heights(&self) -> impl DoubleEndedIterator<Item = u32> + '_ {
+    pub fn relevant_heights(&self) -> impl Iterator<Item = u32> + '_ {
         let mut last_height = Option::<u32>::None;
         self.anchors
             .iter()
@@ -944,17 +935,22 @@ impl<A> Additions<A> {
             })
             .chain(self.txout.iter().map(|(op, txout)| (*op, txout)))
     }
+}
 
-    /// Appends the changes in `other` into self such that applying `self` afterward has the same
-    /// effect as sequentially applying the original `self` and `other`.
-    pub fn append(&mut self, mut other: Additions<A>)
-    where
-        A: Ord,
-    {
+impl<A: Ord> Append for Additions<A> {
+    fn append(&mut self, mut other: Self) {
         self.tx.append(&mut other.tx);
         self.txout.append(&mut other.txout);
         self.anchors.append(&mut other.anchors);
-        self.last_seen.append(&mut other.last_seen);
+
+        // last_seen timestamps should only increase
+        self.last_seen.extend(
+            other
+                .last_seen
+                .into_iter()
+                .filter(|(txid, update_ls)| self.last_seen.get(txid) < Some(update_ls))
+                .collect::<Vec<_>>(),
+        );
     }
 }
 
diff --git a/crates/chain/tests/test_tx_graph.rs b/crates/chain/tests/test_tx_graph.rs
@@ -4,7 +4,7 @@ use bdk_chain::{
     collections::*,
     local_chain::LocalChain,
     tx_graph::{Additions, TxGraph},
-    BlockId, ObservedAs,
+    Append, BlockId, ObservedAs,
 };
 use bitcoin::{
     hashes::Hash, BlockHash, OutPoint, PackedLockTime, Script, Transaction, TxIn, TxOut, Txid,
@@ -849,3 +849,34 @@ fn test_relevant_heights() {
         "anchor for non-existant tx is inserted at height 5, must still be in relevant heights",
     );
 }
+
+/// Ensure that `last_seen` values only increase during [`Append::append`].
+#[test]
+fn test_additions_last_seen_append() {
+    let txid: Txid = h!("test txid");
+
+    let test_cases: &[(Option<u64>, Option<u64>)] = &[
+        (Some(5), Some(6)),
+        (Some(5), Some(5)),
+        (Some(6), Some(5)),
+        (None, Some(5)),
+        (Some(5), None),
+    ];
+
+    for (original_ls, update_ls) in test_cases {
+        let mut original = Additions::<()> {
+            last_seen: original_ls.map(|ls| (txid, ls)).into_iter().collect(),
+            ..Default::default()
+        };
+        let update = Additions::<()> {
+            last_seen: update_ls.map(|ls| (txid, ls)).into_iter().collect(),
+            ..Default::default()
+        };
+
+        original.append(update);
+        assert_eq!(
+            &original.last_seen.get(&txid).cloned(),
+            Ord::max(original_ls, update_ls),
+        );
+    }
+}
