feat(chain): add map_anchors for TxGraph

crates/chain/src/tx_graph.rs

@@ -449,6 +449,32 @@ impl<A> TxGraph<A> {
     pub fn is_empty(&self) -> bool {
         self.txs.is_empty()
     }
+
+    /// Transform the [`TxGraph`] to have [`Anchor`]s of another type.
+    ///
+    /// This takes in a closure of signature `FnOnce(A) -> A2` which is called for each [`Anchor`] to
+    /// transform it.
+    pub fn map_anchors<A2: Ord, F>(self, f: F) -> TxGraph<A2>
+    where
+        F: FnOnce(A) -> A2 + Copy,
+    {
+        TxGraph::<A2> {
+            txs: HashMap::from_iter(self.txs.into_iter().map(
+                |(txid, (tx_node_interal, set, amt))| {
+                    let mut new_set = BTreeSet::<A2>::new();
+                    for anchor in set.into_iter() {
+                        new_set.insert(f(anchor));
+                    }
+                    (txid, (tx_node_interal, new_set, amt))
+                },
+            )),
+            anchors: BTreeSet::<(A2, Txid)>::from_iter(
+                self.anchors.into_iter().map(|(a, txid)| (f(a), txid)),
+            ),
+            spends: self.spends,
+            empty_outspends: self.empty_outspends,
+        }
+    }
 }
 
 impl<A: Clone + Ord> TxGraph<A> {

crates/chain/tests/test_tx_graph.rs

@@ -1172,3 +1172,133 @@ fn test_missing_blocks() {
         ),
     ]);
 }
+
+#[test]
+fn map_tx_graph_to_confirm() {
+    // 2 (Outpoint, TxOut) tuples that denotes original data in the graph, as partial transactions.
+    let original_ops = [
+        (
+            OutPoint::new(h!("tx1"), 1),
+            TxOut {
+                value: 10_000,
+                script_pubkey: ScriptBuf::new(),
+            },
+        ),
+        (
+            OutPoint::new(h!("tx1"), 2),
+            TxOut {
+                value: 20_000,
+                script_pubkey: ScriptBuf::new(),
+            },
+        ),
+    ];
+
+    // Another (OutPoint, TxOut) tuple to be used as update as partial transaction.
+    let update_ops = [(
+        OutPoint::new(h!("tx2"), 0),
+        TxOut {
+            value: 20_000,
+            script_pubkey: ScriptBuf::new(),
+        },
+    )];
+
+    // One full transaction to be included in the update
+    let update_txs = Transaction {
+        version: 0x01,
+        lock_time: absolute::LockTime::ZERO,
+        input: vec![TxIn {
+            previous_output: OutPoint::null(),
+            ..Default::default()
+        }],
+        output: vec![TxOut {
+            value: 30_000,
+            script_pubkey: ScriptBuf::new(),
+        }],
+    };
+
+    // Conf anchor used to mark the full transaction as confirmed.
+    let conf_anchor = ChainPosition::Confirmed(BlockId {
+        height: 100,
+        hash: h!("random blockhash"),
+    });
+
+    // Unconfirmed anchor to mark the partial transactions as unconfirmed
+    let unconf_anchor = ChainPosition::<BlockId>::Unconfirmed(1000000);
+
+    // Make the original graph
+    let mut graph = {
+        let mut graph = TxGraph::<ChainPosition<BlockId>>::default();
+        for (outpoint, txout) in &original_ops {
+            assert_eq!(
+                graph.insert_txout(*outpoint, txout.clone()),
+                ChangeSet {
+                    txouts: [(*outpoint, txout.clone())].into(),
+                    ..Default::default()
+                }
+            );
+        }
+        graph
+    };
+
+    // Add an unconfirmed block to the update graph
+    let update = {
+        let mut graph = TxGraph::default();
+        for (outpoint, txout) in &update_ops {
+            // Insert partials transactions
+            assert_eq!(
+                graph.insert_txout(*outpoint, txout.clone()),
+                ChangeSet {
+                    txouts: [(*outpoint, txout.clone())].into(),
+                    ..Default::default()
+                }
+            );
+            // Mark them unconfirmed.
+            assert_eq!(
+                graph.insert_anchor(outpoint.txid, unconf_anchor),
+                ChangeSet {
+                    txs: [].into(),
+                    txouts: [].into(),
+                    anchors: [(unconf_anchor, outpoint.txid)].into(),
+                    last_seen: [].into()
+                }
+            );
+            // Mark them last seen at.
+            assert_eq!(
+                graph.insert_seen_at(outpoint.txid, 1000000),
+                ChangeSet {
+                    txs: [].into(),
+                    txouts: [].into(),
+                    anchors: [].into(),
+                    last_seen: [(outpoint.txid, 1000000)].into()
+                }
+            );
+        }
+        // Insert the full transaction
+        assert_eq!(
+            graph.insert_tx(update_txs.clone()),
+            ChangeSet {
+                txs: [update_txs.clone()].into(),
+                ..Default::default()
+            }
+        );
+        // No confirm here
+        graph
+    };
+
+    // Check the resulting addition.
+    let _ = graph.apply_update(update);
+
+    // Confirm blocks by map_anchors function
+    let mapped_graph = graph.map_anchors(|anchor| {
+        if let ChainPosition::<BlockId>::Unconfirmed(ts) = anchor {
+            if ts == 1000000 {
+                return conf_anchor;
+            }
+        }
+        anchor
+    });
+
+    let anchors = mapped_graph.all_anchors();
+    assert_eq!(anchors.len(), 1);
+    assert_eq!(anchors.iter().next().unwrap().0, conf_anchor);
+}