feat!: Introduce BroadcastQueue

Author: evanlinjin

wallet/src/types.rs

@@ -73,6 +73,8 @@ pub struct LocalOutput {
     pub derivation_index: u32,
     /// The position of the output in the blockchain.
     pub chain_position: ChainPosition<ConfirmationBlockTime>,
+    /// Whether this output exists in a transaction that is yet to be broadcasted.
+    pub needs_broadcast: bool,
 }
 
 /// A [`Utxo`] with its `satisfaction_weight`.

wallet/src/wallet/broadcast_queue.rs

@@ -1,14 +1,15 @@
 //! Unbroadcasted transaction queue.
 
-use std::vec::Vec;
+use alloc::vec::Vec;
+use chain::tx_graph;
+use chain::Anchor;
+use chain::TxGraph;
 
-use crate::collections::BTreeMap;
-use crate::collections::HashMap;
+use crate::collections::HashSet;
 use crate::collections::VecDeque;
 
 use bitcoin::Txid;
 use chain::Merge;
-use chain::TxUpdate;
 
 /// An ordered unbroadcasted list.
 ///
@@ -17,46 +18,34 @@ use chain::TxUpdate;
 pub struct BroadcastQueue {
     queue: VecDeque<Txid>,
 
-    /// Enforces that we only change the tx entry if the change has a larger `seq`.
-    latest_seq: HashMap<Txid, (u64, bool)>,
-    next_seq: u64,
+    /// Enforces that we do not have duplicates in `queue`.
+    dedup: HashSet<Txid>,
 }
 
-/// Represents changes made to [`BroadcastQueue`].
+/// Represents a single mutation to [`BroadcastQueue`].
+#[derive(Debug, Clone, PartialEq, serde::Deserialize, serde::Serialize)]
+pub enum Mutation {
+    /// A push to the back of the queue.
+    Push(Txid),
+    ///
+    Remove(Txid),
+}
+
+/// A list of mutations made to [`BroadcastQueue`].
 #[must_use]
 #[derive(Debug, Clone, Default, PartialEq, serde::Deserialize, serde::Serialize)]
 pub struct ChangeSet {
-    pub insert_tx: BTreeMap<Txid, u64>,
-    pub remove_tx: BTreeMap<Txid, u64>,
+    /// Mutations.
+    pub mutations: Vec<Mutation>,
 }
 
 impl Merge for ChangeSet {
     fn merge(&mut self, other: Self) {
-        self.insert_tx.extend(
-            other
-                .insert_tx
-                .into_iter()
-                // Only replace `seq`uences which are larger.
-                .filter(|(txid, other_seq)| match self.insert_tx.get(txid) {
-                    Some(seq) => seq < other_seq,
-                    None => true,
-                })
-                .collect::<Vec<_>>(),
-        );
-        self.remove_tx.extend(
-            other
-                .remove_tx
-                .into_iter()
-                .filter(|(txid, other_seq)| match self.remove_tx.get(txid) {
-                    Some(seq) => seq < other_seq,
-                    None => true,
-                })
-                .collect::<Vec<_>>(),
-        );
+        self.mutations.extend(other.mutations);
     }
 
     fn is_empty(&self) -> bool {
-        self.insert_tx.is_empty() && self.remove_tx.is_empty()
+        self.mutations.is_empty()
     }
 }
 
@@ -70,119 +59,175 @@ impl BroadcastQueue {
 
     /// Apply the given `changeset`.
     pub fn apply_changeset(&mut self, changeset: ChangeSet) {
-        for (txid, seq) in changeset.insert_tx {
-            self._push(txid, seq);
-        }
-        for (txid, seq) in changeset.remove_tx {
-            self._remove(txid, seq);
-        }
-        self._sort_queue();
-    }
-
-    fn _sort_queue(&mut self) {
-        self.queue
-            .make_contiguous()
-            .sort_unstable_by_key(|txid| self.latest_seq.get(txid));
-    }
-
-    fn _push(&mut self, txid: Txid, seq: u64) -> bool {
-        let modified = &mut false;
-        self.latest_seq
-            .entry(txid)
-            .and_modify(|(v, in_queue)| {
-                if *v < seq {
-                    *v = seq;
-                    *modified = true;
-                    if !*in_queue {
-                        *in_queue = true;
-                        self.queue.push_back(txid);
-                    }
-                }
-            })
-            .or_insert_with(|| {
-                *modified = true;
-                self.queue.push_back(txid);
-                (seq, true)
-            });
-        self.next_seq = self.next_seq.max(seq.saturating_add(1));
-        *modified
-    }
-
-    fn _contains(latest_seq: &mut HashMap<Txid, (u64, bool)>, txid: Txid) -> bool {
-        match latest_seq.get(&txid) {
-            Some(&(_, contains)) => contains,
-            None => false,
+        for mutation in changeset.mutations {
+            match mutation {
+                Mutation::Push(txid) => self._push(txid),
+                Mutation::Remove(txid) => self._remove(txid),
+            };
         }
     }
 
-    fn _remove(&mut self, txid: Txid, seq: u64) -> bool {
-        let modified = &mut false;
-        self.latest_seq
-            .entry(txid)
-            .and_modify(|(v, in_queue)| {
-                if *v <= seq {
-                    *v = seq;
-                    *modified = true;
-                    if *in_queue {
-                        *in_queue = false;
-                        // Most of the time, we are popping from the front of the queue.
-                        let index = self
-                            .queue
-                            .iter()
-                            .enumerate()
-                            .find_map(|(i, &q_txid)| if q_txid == txid { Some(i) } else { None })
-                            .expect("must be in queue");
-                        self.queue.swap_remove_front(index);
-                    }
-                }
-            })
-            .or_insert_with(|| {
-                *modified = true;
-                (seq, false)
-            });
-        self.next_seq = self.next_seq.max(seq.saturating_add(1));
-        *modified
+    /// Whether the `txid` exists in the queue.
+    pub fn contains(&self, txid: Txid) -> bool {
+        self.dedup.contains(&txid)
     }
 
-    /// Reinserting will bump the tx's seq to `next_seq`.
+    /// Push a `txid` to the queue if it does not already exist.
+    ///
+    /// # Warning
+    ///
+    /// This does not get rid of conflicting transactions already in the queue.
     pub fn push(&mut self, txid: Txid) -> ChangeSet {
-        let seq = self.next_seq;
+        let mut changeset = ChangeSet::default();
+        if self._push(txid) {
+            changeset.mutations.push(Mutation::Push(txid));
+        }
+        changeset
+    }
+    fn _push(&mut self, txid: Txid) -> bool {
+        if self.dedup.insert(txid) {
+            self.queue.push_back(txid);
+            return true;
+        }
+        false
+    }
 
+    /// Push a `txid` to the broadcast queue (if it does not exist already) and displaces all
+    /// coflicting txids in the queue.
+    pub fn push_and_displace_conflicts<A>(&mut self, tx_graph: &TxGraph<A>, txid: Txid) -> ChangeSet
+    where
+        A: Anchor,
+    {
         let mut changeset = ChangeSet::default();
-        if self._push(txid, seq) {
-            changeset.insert_tx.insert(txid, seq);
-            self.next_seq += 1;
+
+        let tx = match tx_graph.get_tx(txid) {
+            Some(tx) => tx,
+            None => {
+                debug_assert!(
+                    !self.dedup.contains(&txid),
+                    "Cannot have txid in queue which has no corresponding tx in graph"
+                );
+                return changeset;
+            }
+        };
+
+        if self._push(txid) {
+            changeset.mutations.push(Mutation::Push(txid));
+
+            for txid in tx_graph.walk_conflicts(&tx, |_, conflict_txid| Some(conflict_txid)) {
+                if self._remove(txid) {
+                    changeset.mutations.push(Mutation::Remove(txid));
+                }
+            }
         }
+
         changeset
     }
 
-    /// Returns the next `txid` of the transaction to broadcast.
-    pub fn next(&self) -> Option<Txid> {
-        // TODO: Implement.
-        // NOTE: Prioritize txs that have no dependencies.
-        None
+    /// Returns the next `txid` of the queue to broadcast which has no dependencies to other
+    /// transactions in the queue.
+    pub fn next_to_broadcast<A>(&self, tx_graph: &TxGraph<A>) -> Option<Txid>
+    where
+        A: Anchor,
+    {
+        self.queue.iter().copied().find(|&txid| {
+            let tx = match tx_graph.get_tx(txid) {
+                Some(tx) => tx,
+                None => return false,
+            };
+            if tx
+                .input
+                .iter()
+                .any(|txin| self.dedup.contains(&txin.previous_output.txid))
+            {
+                return false;
+            }
+            true
+        })
     }
 
-    /// Untrack the `txid`.
-    pub fn remove(&mut self, txid: Txid) -> ChangeSet {
-        let seq = self.next_seq;
+    /// Returns unbroadcasted dependencies of the given `txid`.
+    ///
+    /// The returned `Vec` is in broadcast order.
+    pub fn unbroadcasted_dependencies<A>(&self, tx_graph: &TxGraph<A>, txid: Txid) -> Vec<Txid>
+    where
+        A: Anchor,
+    {
+        let tx = match tx_graph.get_tx(txid) {
+            Some(tx) => tx,
+            None => return Vec::new(),
+        };
+        let mut txs = tx_graph
+            .walk_ancestors(tx, |_depth, ancestor_tx| {
+                let ancestor_txid = ancestor_tx.compute_txid();
+                if self.dedup.contains(&ancestor_txid) {
+                    Some(ancestor_txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+        txs.reverse();
+        txs
+    }
 
+    /// Untracks and removes a transaction from the broadcast queue.
+    ///
+    /// Transactions are automatically removed from the queue upon successful broadcast, so calling
+    /// this method directly is typically not required.
+    pub fn remove(&mut self, txid: Txid) -> ChangeSet {
         let mut changeset = ChangeSet::default();
-        if self._remove(txid, seq) {
-            changeset.remove_tx.insert(txid, seq);
-            self.next_seq += 1;
+        if self._remove(txid) {
+            changeset.mutations.push(Mutation::Remove(txid));
         }
         changeset
     }
+    fn _remove(&mut self, txid: Txid) -> bool {
+        if self.dedup.remove(&txid) {
+            let i = (0..self.queue.len())
+                .zip(self.queue.iter().copied())
+                .find_map(|(i, queue_txid)| if queue_txid == txid { Some(i) } else { None })
+                .expect("must exist in queue to exist in `queue`");
+            let _removed = self.queue.remove(i);
+            debug_assert_eq!(_removed, Some(txid));
+            return true;
+        }
+        false
+    }
 
-    /// Untrack transactions that are given anchors and seen-at timestamps.
-    pub fn update<A>(&mut self, tx_update: &TxUpdate<A>) -> ChangeSet {
+    /// Untracks and removes a transaction and it's descendants from the broadcast queue.
+    pub fn remove_and_displace_dependants<A>(
+        &mut self,
+        tx_graph: &TxGraph<A>,
+        txid: Txid,
+    ) -> ChangeSet
+    where
+        A: Anchor,
+    {
         let mut changeset = ChangeSet::default();
-        for (_, txid) in &tx_update.anchors {
-            changeset.merge(self.remove(*txid));
+
+        if self._remove(txid) {
+            changeset.mutations.push(Mutation::Remove(txid));
+            for txid in tx_graph.walk_descendants(txid, |_depth, txid| Some(txid)) {
+                if self._remove(txid) {
+                    changeset.mutations.push(Mutation::Remove(txid));
+                }
+            }
         }
-        for (txid, _) in &tx_update.seen_ats {
-            changeset.merge(self.remove(*txid));
+        changeset
+    }
+
+    /// Untrack transactions that are given anchors and/or mempool timestamps.
+    pub fn filter_from_graph_changeset<A>(
+        &mut self,
+        graph_changeset: &tx_graph::ChangeSet<A>,
+    ) -> ChangeSet {
+        let mut changeset = ChangeSet::default();
+        let s_txids = graph_changeset.last_seen.iter().map(|(txid, _)| *txid);
+        let a_txids = graph_changeset.anchors.iter().map(|(_, txid)| *txid);
+        let e_txids = graph_changeset.last_evicted.iter().map(|(txid, _)| *txid);
+        for txid in s_txids.chain(a_txids).chain(e_txids) {
+            changeset.merge(self.remove(txid));
         }
         changeset
     }
@@ -193,4 +238,11 @@ impl BroadcastQueue {
     pub fn txids(&self) -> impl ExactSizeIterator<Item = Txid> + '_ {
         self.queue.iter().copied()
     }
+
+    /// Initial changeset.
+    pub fn initial_changeset(&self) -> ChangeSet {
+        ChangeSet {
+            mutations: self.queue.iter().copied().map(Mutation::Push).collect(),
+        }
+    }
 }