kaspa_mining/model/
topological_sort.rs

1use itertools::Itertools;
2use kaspa_consensus_core::tx::Transaction;
3use std::{
4    collections::{HashMap, HashSet, VecDeque},
5    iter::{FusedIterator, Map},
6};
7
8type IndexSet = HashSet<usize>;
9
10pub trait TopologicalSort {
11    fn topological_sort(self) -> Self
12    where
13        Self: Sized;
14}
15
16impl<T: AsRef<Transaction> + Clone> TopologicalSort for Vec<T> {
17    fn topological_sort(self) -> Self {
18        let mut sorted = Vec::with_capacity(self.len());
19        let mut in_degree: Vec<i32> = vec![0; self.len()];
20
21        // Index on transaction ids
22        let mut index = HashMap::with_capacity(self.len());
23        self.iter().enumerate().for_each(|(idx, tx)| {
24            let _ = index.insert(tx.as_ref().id(), idx);
25        });
26
27        // Transaction edges
28        let mut all_edges: Vec<Option<IndexSet>> = vec![None; self.len()];
29        self.iter().enumerate().for_each(|(destination_idx, tx)| {
30            tx.as_ref().inputs.iter().for_each(|input| {
31                if let Some(origin_idx) = index.get(&input.previous_outpoint.transaction_id) {
32                    all_edges[*origin_idx].get_or_insert_with(IndexSet::new).insert(destination_idx);
33                }
34            })
35        });
36
37        // Degrees
38        (0..self.len()).for_each(|origin_idx| {
39            if let Some(ref edges) = all_edges[origin_idx] {
40                edges.iter().for_each(|destination_idx| {
41                    in_degree[*destination_idx] += 1;
42                });
43            }
44        });
45
46        // Degree 0
47        let mut queue = VecDeque::with_capacity(self.len());
48        (0..self.len()).for_each(|destination_idx| {
49            if in_degree[destination_idx] == 0 {
50                queue.push_back(destination_idx);
51            }
52        });
53
54        // Sorted transactions
55        while !queue.is_empty() {
56            let current = queue.pop_front().unwrap();
57            if let Some(ref edges) = all_edges[current] {
58                edges.iter().for_each(|destination_idx| {
59                    let degree = in_degree.get_mut(*destination_idx).unwrap();
60                    *degree -= 1;
61                    if *degree == 0 {
62                        queue.push_back(*destination_idx);
63                    }
64                });
65            }
66            sorted.push(self[current].clone());
67        }
68        assert_eq!(sorted.len(), self.len(), "by definition, cryptographically no cycle can exist in a DAG of transactions");
69
70        sorted
71    }
72}
73
74pub trait IterTopologically<T>
75where
76    T: AsRef<Transaction>,
77{
78    fn topological_iter(&self) -> TopologicalIter<'_, T>;
79}
80
81impl<T: AsRef<Transaction>> IterTopologically<T> for &[T] {
82    fn topological_iter(&self) -> TopologicalIter<'_, T> {
83        TopologicalIter::new(self)
84    }
85}
86
87impl<T: AsRef<Transaction>> IterTopologically<T> for Vec<T> {
88    fn topological_iter(&self) -> TopologicalIter<'_, T> {
89        TopologicalIter::new(self)
90    }
91}
92
93pub struct TopologicalIter<'a, T: AsRef<Transaction>> {
94    transactions: &'a [T],
95    in_degree: Vec<i32>,
96    edges: Vec<Option<IndexSet>>,
97    queue: VecDeque<usize>,
98    yields_count: usize,
99}
100
101impl<'a, T: AsRef<Transaction>> TopologicalIter<'a, T> {
102    pub fn new(transactions: &'a [T]) -> Self {
103        let mut in_degree: Vec<i32> = vec![0; transactions.len()];
104
105        // Index on transaction ids
106        let mut index = HashMap::with_capacity(transactions.len());
107        transactions.iter().enumerate().for_each(|(idx, tx)| {
108            let _ = index.insert(tx.as_ref().id(), idx);
109        });
110
111        // Transaction edges
112        let mut edges: Vec<Option<IndexSet>> = vec![None; transactions.len()];
113        transactions.iter().enumerate().for_each(|(destination_idx, tx)| {
114            tx.as_ref().inputs.iter().for_each(|input| {
115                if let Some(origin_idx) = index.get(&input.previous_outpoint.transaction_id) {
116                    edges[*origin_idx].get_or_insert_with(IndexSet::new).insert(destination_idx);
117                }
118            })
119        });
120
121        // Degrees
122        (0..transactions.len()).for_each(|origin_idx| {
123            if let Some(ref edges) = edges[origin_idx] {
124                edges.iter().for_each(|destination_idx| {
125                    in_degree[*destination_idx] += 1;
126                });
127            }
128        });
129
130        // Degree 0
131        let mut queue = VecDeque::with_capacity(transactions.len());
132        (0..transactions.len()).for_each(|destination_idx| {
133            if in_degree[destination_idx] == 0 {
134                queue.push_back(destination_idx);
135            }
136        });
137        Self { transactions, in_degree, edges, queue, yields_count: 0 }
138    }
139}
140
141impl<'a, T: AsRef<Transaction>> Iterator for TopologicalIter<'a, T> {
142    type Item = &'a T;
143
144    fn next(&mut self) -> Option<Self::Item> {
145        match self.queue.pop_front() {
146            Some(current) => {
147                if let Some(ref edges) = self.edges[current] {
148                    edges.iter().for_each(|destination_idx| {
149                        let degree = self.in_degree.get_mut(*destination_idx).unwrap();
150                        *degree -= 1;
151                        if *degree == 0 {
152                            self.queue.push_back(*destination_idx);
153                        }
154                    });
155                }
156                self.yields_count += 1;
157                Some(&self.transactions[current])
158            }
159            None => None,
160        }
161    }
162
163    fn size_hint(&self) -> (usize, Option<usize>) {
164        let items_remaining = self.transactions.len() - self.yields_count.min(self.transactions.len());
165        (self.yields_count, Some(items_remaining))
166    }
167}
168
169impl<'a, T: AsRef<Transaction>> FusedIterator for TopologicalIter<'a, T> {}
170impl<'a, T: AsRef<Transaction>> ExactSizeIterator for TopologicalIter<'a, T> {
171    fn len(&self) -> usize {
172        self.transactions.len()
173    }
174}
175
176pub trait IntoIterTopologically<T>
177where
178    T: AsRef<Transaction>,
179{
180    fn topological_into_iter(self) -> TopologicalIntoIter<T>;
181}
182
183impl<T: AsRef<Transaction>> IntoIterTopologically<T> for Vec<T> {
184    fn topological_into_iter(self) -> TopologicalIntoIter<T> {
185        TopologicalIntoIter::new(self)
186    }
187}
188
189impl<T, I, F> IntoIterTopologically<T> for Map<I, F>
190where
191    T: AsRef<Transaction>,
192    I: Iterator,
193    F: FnMut(<I as Iterator>::Item) -> T,
194{
195    fn topological_into_iter(self) -> TopologicalIntoIter<T> {
196        TopologicalIntoIter::new(self)
197    }
198}
199
200pub struct TopologicalIntoIter<T: AsRef<Transaction>> {
201    transactions: Vec<Option<T>>,
202    in_degree: Vec<i32>,
203    edges: Vec<Option<IndexSet>>,
204    queue: VecDeque<usize>,
205    yields_count: usize,
206}
207
208impl<T: AsRef<Transaction>> TopologicalIntoIter<T> {
209    pub fn new(transactions: impl IntoIterator<Item = T>) -> Self {
210        // Collect all transactions
211        let transactions = transactions.into_iter().map(|tx| Some(tx)).collect_vec();
212
213        let mut in_degree: Vec<i32> = vec![0; transactions.len()];
214
215        // Index on transaction ids
216        let mut index = HashMap::with_capacity(transactions.len());
217        transactions.iter().enumerate().for_each(|(idx, tx)| {
218            let _ = index.insert(tx.as_ref().unwrap().as_ref().id(), idx);
219        });
220
221        // Transaction edges
222        let mut edges: Vec<Option<IndexSet>> = vec![None; transactions.len()];
223        transactions.iter().enumerate().for_each(|(destination_idx, tx)| {
224            tx.as_ref().unwrap().as_ref().inputs.iter().for_each(|input| {
225                if let Some(origin_idx) = index.get(&input.previous_outpoint.transaction_id) {
226                    edges[*origin_idx].get_or_insert_with(IndexSet::new).insert(destination_idx);
227                }
228            })
229        });
230
231        // Degrees
232        (0..transactions.len()).for_each(|origin_idx| {
233            if let Some(ref edges) = edges[origin_idx] {
234                edges.iter().for_each(|destination_idx| {
235                    in_degree[*destination_idx] += 1;
236                });
237            }
238        });
239
240        // Degree 0
241        let mut queue = VecDeque::with_capacity(transactions.len());
242        (0..transactions.len()).for_each(|destination_idx| {
243            if in_degree[destination_idx] == 0 {
244                queue.push_back(destination_idx);
245            }
246        });
247        Self { transactions, in_degree, edges, queue, yields_count: 0 }
248    }
249}
250
251impl<T: AsRef<Transaction>> Iterator for TopologicalIntoIter<T> {
252    type Item = T;
253
254    fn next(&mut self) -> Option<Self::Item> {
255        match self.queue.pop_front() {
256            Some(current) => {
257                if let Some(ref edges) = self.edges[current] {
258                    edges.iter().for_each(|destination_idx| {
259                        let degree = self.in_degree.get_mut(*destination_idx).unwrap();
260                        *degree -= 1;
261                        if *degree == 0 {
262                            self.queue.push_back(*destination_idx);
263                        }
264                    });
265                }
266                self.yields_count += 1;
267                self.transactions[current].take()
268            }
269            None => None,
270        }
271    }
272
273    fn size_hint(&self) -> (usize, Option<usize>) {
274        let items_remaining = self.transactions.len() - self.yields_count.min(self.transactions.len());
275        (self.yields_count, Some(items_remaining))
276    }
277}
278
279impl<T: AsRef<Transaction>> FusedIterator for TopologicalIntoIter<T> {}
280impl<T: AsRef<Transaction>> ExactSizeIterator for TopologicalIntoIter<T> {
281    fn len(&self) -> usize {
282        self.transactions.len()
283    }
284}