ark_core/
tx_graph.rs

1use crate::Error;
2use crate::VTXO_INPUT_INDEX;
3use bitcoin::taproot::Signature;
4use bitcoin::Psbt;
5use bitcoin::Txid;
6use std::collections::HashMap;
7
8#[derive(Debug, Clone)]
9pub struct TxGraph {
10    root: Psbt,
11    children: HashMap<u32, TxGraph>,
12}
13
14#[derive(Debug, Clone)]
15pub struct TxGraphChunk {
16    pub txid: Option<Txid>,
17    pub tx: Psbt,
18    pub children: HashMap<u32, Txid>,
19}
20
21impl TxGraph {
22    pub fn new(chunks: Vec<TxGraphChunk>) -> Result<Self, Error> {
23        if chunks.is_empty() {
24            return Err(Error::ad_hoc("empty chunks"));
25        }
26
27        // Create a map to store all chunks by their txid for easy lookup.
28        let mut chunks_by_txid = HashMap::new();
29
30        for chunk in chunks {
31            let txid = chunk.tx.unsigned_tx.compute_txid();
32            chunks_by_txid.insert(txid, chunk);
33        }
34
35        // Find the root chunks (the ones that aren't referenced as a child).
36        let mut root_txids = Vec::new();
37
38        for txid in chunks_by_txid.keys() {
39            let mut is_child = false;
40
41            for (other_txid, other_chunk) in &chunks_by_txid {
42                if other_txid == txid {
43                    // Skip self.
44                    continue;
45                }
46
47                // Check if the current chunk is a child of the other chunk.
48                if other_chunk
49                    .children
50                    .values()
51                    .any(|child_txid| child_txid == txid)
52                {
53                    is_child = true;
54                    break;
55                }
56            }
57
58            // If the chunk is not a child of any other chunk, it is a root.
59            if !is_child {
60                root_txids.push(*txid);
61            }
62        }
63
64        if root_txids.is_empty() {
65            return Err(Error::ad_hoc("no root chunk found"));
66        }
67
68        if root_txids.len() > 1 {
69            return Err(Error::ad_hoc(format!(
70                "multiple root chunks found: {root_txids:?}",
71            )));
72        }
73
74        let graph = Self::build_graph(root_txids[0], &chunks_by_txid).ok_or_else(|| {
75            Error::ad_hoc(format!("chunk not found for root txid: {}", root_txids[0]))
76        })?;
77
78        // Verify that the number of chunks is equal to the number of nodes in the graph
79        let chunk_count = chunks_by_txid.len();
80        let node_count = graph.nb_of_nodes();
81        if node_count != chunk_count {
82            return Err(Error::ad_hoc(format!(
83                "number of chunks ({chunk_count}) is not equal to \
84                 the number of nodes in the graph ({node_count})",
85            )));
86        }
87
88        Ok(graph)
89    }
90
91    fn build_graph(root_txid: Txid, chunks_by_txid: &HashMap<Txid, TxGraphChunk>) -> Option<Self> {
92        let root_chunk = chunks_by_txid.get(&root_txid)?;
93
94        let mut children = HashMap::new();
95
96        for (output_index, child_txid) in &root_chunk.children {
97            if let Some(child_graph) = Self::build_graph(*child_txid, chunks_by_txid) {
98                children.insert(*output_index, child_graph);
99            }
100        }
101
102        Some(TxGraph {
103            root: root_chunk.tx.clone(),
104            children,
105        })
106    }
107
108    pub fn nb_of_nodes(&self) -> usize {
109        let mut nb = 1;
110        for child in self.children.values() {
111            nb += child.nb_of_nodes();
112        }
113        nb
114    }
115
116    pub fn apply<F>(&mut self, f: F) -> Result<(), Error>
117    where
118        F: Fn(&mut TxGraph) -> Result<bool, Error> + Copy,
119    {
120        let should_continue = f(self)?;
121
122        if !should_continue {
123            return Ok(());
124        }
125
126        for child in self.children.values_mut() {
127            child.apply(f)?;
128        }
129
130        Ok(())
131    }
132
133    pub fn find(&self, txid: &Txid) -> Option<&Self> {
134        if self.root.unsigned_tx.compute_txid() == *txid {
135            return Some(self);
136        }
137
138        for child in self.children.values() {
139            if let Some(node) = child.find(txid) {
140                return Some(node);
141            }
142        }
143
144        None
145    }
146
147    pub fn as_map(&self) -> HashMap<Txid, &Psbt> {
148        fn _as_map<'a>(graph: &'a TxGraph, map: &mut HashMap<Txid, &'a Psbt>) {
149            map.insert(graph.root.unsigned_tx.compute_txid(), &graph.root);
150
151            for (_, child) in graph.children.iter() {
152                _as_map(child, map);
153            }
154        }
155
156        let mut map = HashMap::new();
157        _as_map(self, &mut map);
158
159        map
160    }
161
162    pub fn set_signature(&mut self, sig: Signature) {
163        self.root.inputs[VTXO_INPUT_INDEX].tap_key_sig = Some(sig);
164    }
165
166    pub fn root(&self) -> &Psbt {
167        &self.root
168    }
169
170    /// Return all leaf nodes (transactions without children) in the graph.
171    pub fn leaves(&self) -> Vec<&Psbt> {
172        if self.children.is_empty() {
173            return vec![&self.root];
174        }
175
176        let mut leaves = Vec::new();
177        for child in self.children.values() {
178            leaves.extend(child.leaves());
179        }
180
181        leaves
182    }
183}
ark_core/tx_graph.rs

ark_core/
tx_graph.rs
