feat(chain): Signed txs should displace unsigned txs in TxGraph.

Author: evanlinjin

crates/chain/src/tx_graph.rs

@@ -613,22 +613,79 @@ impl<A: Anchor> TxGraph<A> {
         changeset
     }
 
-    /// Inserts the given transaction into [`TxGraph`].
+    /// Insert the given transaction into [`TxGraph`].
     ///
-    /// The [`ChangeSet`] returned will be empty if `tx` already exists.
+    /// The [`ChangeSet`] returned will be empty if no changes are made to the graph.
+    ///
+    /// # Updating Existing Transactions
+    ///
+    /// An unsigned transaction can be inserted first and have it's witness fields updated with
+    /// further transaction insertions (given that the newly introduced transaction shares the same
+    /// txid as the original transaction).
+    ///
+    /// The witnesses of the newly introduced transaction will be merged with the witnesses of the
+    /// original transaction in a way where:
+    ///
+    /// * A non-empty witness has precedence over an empty witness.
+    /// * A smaller witness has precedence over a larger witness.
+    /// * If the witness sizes are the same, we prioritize the two witnesses with lexicographical
+    ///     order.
     pub fn insert_tx<T: Into<Arc<Transaction>>>(&mut self, tx: T) -> ChangeSet<A> {
+        // This returns `Some` only if the merged tx is different to the `original_tx`.
+        fn _merge_tx_witnesses(
+            original_tx: &Arc<Transaction>,
+            other_tx: &Arc<Transaction>,
+        ) -> Option<Arc<Transaction>> {
+            debug_assert_eq!(
+                original_tx.input.len(),
+                other_tx.input.len(),
+                "tx input count must be the same"
+            );
+            let merged_input = Iterator::zip(original_tx.input.iter(), other_tx.input.iter())
+                .map(|(original_txin, other_txin)| {
+                    let original_key = core::cmp::Reverse((
+                        original_txin.witness.is_empty(),
+                        original_txin.witness.size(),
+                        &original_txin.witness,
+                    ));
+                    let other_key = core::cmp::Reverse((
+                        other_txin.witness.is_empty(),
+                        other_txin.witness.size(),
+                        &other_txin.witness,
+                    ));
+                    if original_key > other_key {
+                        original_txin.clone()
+                    } else {
+                        other_txin.clone()
+                    }
+                })
+                .collect::<Vec<_>>();
+            if merged_input == original_tx.input {
+                return None;
+            }
+            if merged_input == other_tx.input {
+                return Some(other_tx.clone());
+            }
+            Some(Arc::new(Transaction {
+                input: merged_input,
+                ..(**original_tx).clone()
+            }))
+        }
+
         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"
-                );
+                if existing_tx.as_ref() != tx.as_ref() {
+                    // Allowing updating witnesses of txs.
+                    if let Some(merged_tx) = _merge_tx_witnesses(existing_tx, &tx) {
+                        *existing_tx = merged_tx.clone();
+                        changeset.txs.insert(merged_tx);
+                    }
+                }
             }
             partial_tx => {
                 for txin in &tx.input {

crates/chain/tests/test_tx_graph.rs

@@ -10,6 +10,8 @@ use bdk_chain::{
     Anchor, ChainOracle, ChainPosition, Merge,
 };
 use bdk_testenv::{block_id, hash, utils::new_tx};
+use bitcoin::hex::FromHex;
+use bitcoin::Witness;
 use bitcoin::{
     absolute, hashes::Hash, transaction, Amount, BlockHash, OutPoint, ScriptBuf, SignedAmount,
     Transaction, TxIn, TxOut, Txid,
@@ -284,6 +286,74 @@ fn insert_tx_displaces_txouts() {
     assert_eq!(tx_graph.get_txout(outpoint), Some(txout));
 }
 
+#[test]
+fn insert_signed_tx_displaces_unsigned() {
+    let previous_output = OutPoint::new(hash!("prev"), 2);
+    let unsigned_tx = Transaction {
+        version: transaction::Version::ONE,
+        lock_time: absolute::LockTime::ZERO,
+        input: vec![TxIn {
+            previous_output,
+            script_sig: ScriptBuf::default(),
+            sequence: transaction::Sequence::ENABLE_RBF_NO_LOCKTIME,
+            witness: Witness::default(),
+        }],
+        output: vec![TxOut {
+            value: Amount::from_sat(24_000),
+            script_pubkey: ScriptBuf::default(),
+        }],
+    };
+    let signed_tx = Transaction {
+        input: vec![TxIn {
+            previous_output,
+            script_sig: ScriptBuf::default(),
+            sequence: transaction::Sequence::ENABLE_RBF_NO_LOCKTIME,
+            witness: Witness::from_slice(&[
+                // Random witness from mempool.space
+                Vec::from_hex("d59118058bf9e8604cec5c0b4a13430b07286482784da313594e932faad074dc4bd27db7cbfff9ad32450db097342d0148ec21c3033b0c27888fd2fd0de2e9b5")
+                    .unwrap(),
+            ]),
+        }],
+        ..unsigned_tx.clone()
+    };
+
+    // Signed tx must displace unsigned.
+    {
+        let mut tx_graph = TxGraph::<ConfirmationBlockTime>::default();
+        let changeset_insert_unsigned = tx_graph.insert_tx(unsigned_tx.clone());
+        let changeset_insert_signed = tx_graph.insert_tx(signed_tx.clone());
+        assert_eq!(
+            changeset_insert_unsigned,
+            ChangeSet {
+                txs: [Arc::new(unsigned_tx.clone())].into(),
+                ..Default::default()
+            }
+        );
+        assert_eq!(
+            changeset_insert_signed,
+            ChangeSet {
+                txs: [Arc::new(signed_tx.clone())].into(),
+                ..Default::default()
+            }
+        );
+    }
+
+    // Unsigned tx must not displace signed.
+    {
+        let mut tx_graph = TxGraph::<ConfirmationBlockTime>::default();
+        let changeset_insert_signed = tx_graph.insert_tx(signed_tx.clone());
+        let changeset_insert_unsigned = tx_graph.insert_tx(unsigned_tx.clone());
+        assert_eq!(
+            changeset_insert_signed,
+            ChangeSet {
+                txs: [Arc::new(signed_tx)].into(),
+                ..Default::default()
+            }
+        );
+        assert!(changeset_insert_unsigned.is_empty());
+    }
+}
+
 #[test]
 fn insert_txout_does_not_displace_tx() {
     let mut tx_graph = TxGraph::<ConfirmationBlockTime>::default();