From c0da72882105ccc9573e657be216b6be04a04c06 Mon Sep 17 00:00:00 2001
From: Shubham Shinde <shindeshubham0520@gmail.com>
Date: Tue, 10 Mar 2026 02:06:35 +0530
Subject: [PATCH] fix(chain): return TxNode from CanonicalIter and simplify
 CanonicalView

---
 crates/chain/src/canonical_iter.rs        | 16 +++---
 crates/chain/src/canonical_view.rs        | 13 ++---
 crates/chain/tests/test_canonical_view.rs | 59 +++++++++++++++++++++++
 3 files changed, 72 insertions(+), 16 deletions(-)

diff --git a/crates/chain/src/canonical_iter.rs b/crates/chain/src/canonical_iter.rs
index 204ead451..7222c1793 100644
--- a/crates/chain/src/canonical_iter.rs
+++ b/crates/chain/src/canonical_iter.rs
@@ -1,5 +1,5 @@
 use crate::collections::{HashMap, HashSet, VecDeque};
-use crate::tx_graph::{TxAncestors, TxDescendants};
+use crate::tx_graph::{TxAncestors, TxDescendants, TxNode};
 use crate::{Anchor, ChainOracle, TxGraph};
 use alloc::boxed::Box;
 use alloc::collections::BTreeSet;
@@ -200,18 +200,22 @@ impl<'g, A: Anchor, C: ChainOracle> CanonicalIter<'g, A, C> {
     }
 }
 
-impl<A: Anchor, C: ChainOracle> Iterator for CanonicalIter<'_, A, C> {
-    type Item = Result<(Txid, Arc<Transaction>, CanonicalReason<A>), C::Error>;
+impl<'g, A: Anchor, C: ChainOracle> Iterator for CanonicalIter<'g, A, C> {
+    type Item = Result<(TxNode<'g, Arc<Transaction>, A>, CanonicalReason<A>), C::Error>;
 
     fn next(&mut self) -> Option<Self::Item> {
         loop {
             if let Some(txid) = self.queue.pop_front() {
-                let (tx, reason) = self
+                let reason = self
                     .canonical
                     .get(&txid)
-                    .cloned()
+                    .map(|(_, reason)| reason.clone())
                     .expect("reason must exist");
-                return Some(Ok((txid, tx, reason)));
+                let tx_node = self
+                    .tx_graph
+                    .get_tx_node(txid)
+                    .expect("tx node must exist if in canonical map");
+                return Some(Ok((tx_node, reason)));
             }
 
             if let Some((txid, tx)) = self.unprocessed_assumed_txs.next() {
diff --git a/crates/chain/src/canonical_view.rs b/crates/chain/src/canonical_view.rs
index 0191f4507..8f8b4d480 100644
--- a/crates/chain/src/canonical_view.rs
+++ b/crates/chain/src/canonical_view.rs
@@ -134,16 +134,9 @@ impl<A: Anchor> CanonicalView<A> {
         };
 
         for r in CanonicalIter::new(tx_graph, chain, chain_tip, params) {
-            let (txid, tx, why) = r?;
-
-            let tx_node = match tx_graph.get_tx_node(txid) {
-                Some(tx_node) => tx_node,
-                None => {
-                    // TODO: Have the `CanonicalIter` return `TxNode`s.
-                    debug_assert!(false, "tx node must exist!");
-                    continue;
-                }
-            };
+            let (tx_node, why) = r?;
+            let txid = tx_node.txid;
+            let tx = tx_node.tx.clone();
 
             view.order.push(txid);
 
diff --git a/crates/chain/tests/test_canonical_view.rs b/crates/chain/tests/test_canonical_view.rs
index 3c0d54381..a629c0a59 100644
--- a/crates/chain/tests/test_canonical_view.rs
+++ b/crates/chain/tests/test_canonical_view.rs
@@ -301,3 +301,62 @@ fn test_min_confirmations_multiple_transactions() {
     );
     assert_eq!(balance_high.untrusted_pending, Amount::ZERO);
 }
+
+#[test]
+fn test_canonical_view_construction_with_tx_nodes() {
+    let blocks: BTreeMap<u32, BlockHash> = [(0, hash!("genesis")), (1, hash!("b1"))]
+        .into_iter()
+        .collect();
+    let chain = LocalChain::from_blocks(blocks).unwrap();
+
+    let mut tx_graph = TxGraph::default();
+    let tx = Transaction {
+        input: vec![TxIn {
+            previous_output: OutPoint::new(hash!("parent"), 0),
+            ..Default::default()
+        }],
+        output: vec![TxOut {
+            value: Amount::from_sat(100_000),
+            script_pubkey: ScriptBuf::new(),
+        }],
+        ..new_tx(1)
+    };
+    let txid = tx.compute_txid();
+    let outpoint = OutPoint::new(txid, 0);
+
+    let _ = tx_graph.insert_tx(tx);
+
+    let _ = tx_graph.insert_anchor(
+        txid,
+        ConfirmationBlockTime {
+            block_id: chain.get(1).unwrap().block_id(),
+            confirmation_time: 123456,
+        },
+    );
+
+    let chain_tip = chain.tip().block_id();
+    let canonical_view =
+        tx_graph.canonical_view(&chain, chain_tip, CanonicalizationParams::default());
+
+    let tx_opt = canonical_view.tx(txid);
+    assert!(
+        tx_opt.is_some(),
+        "Transaction should exist in canonical view"
+    );
+
+    let canonical_tx = tx_opt.unwrap();
+    assert_eq!(canonical_tx.txid, txid);
+
+    let txout = canonical_view.txout(outpoint);
+    assert!(
+        txout.is_some(),
+        "Output should be available in canonical view"
+    );
+
+    let balance = canonical_view.balance(
+        [((), outpoint)],
+        |_, _| true, // trust all
+        1,
+    );
+    assert_eq!(balance.confirmed, Amount::from_sat(100_000));
+}