Skip to main content

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 IterTopologically<T>
11where
12    T: AsRef<Transaction>,
13{
14    fn topological_iter(&self) -> TopologicalIter<'_, T>;
15}
16
17impl<T: AsRef<Transaction>> IterTopologically<T> for &[T] {
18    fn topological_iter(&self) -> TopologicalIter<'_, T> {
19        TopologicalIter::new(self)
20    }
21}
22
23impl<T: AsRef<Transaction>> IterTopologically<T> for Vec<T> {
24    fn topological_iter(&self) -> TopologicalIter<'_, T> {
25        TopologicalIter::new(self)
26    }
27}
28
29pub struct TopologicalIter<'a, T: AsRef<Transaction>> {
30    transactions: &'a [T],
31    in_degree: Vec<isize>,
32    edges: Vec<Option<IndexSet>>,
33    queue: VecDeque<usize>,
34    yields_count: usize,
35}
36
37impl<'a, T: AsRef<Transaction>> TopologicalIter<'a, T> {
38    pub fn new(transactions: &'a [T]) -> Self {
39        let mut in_degree: Vec<isize> = vec![0; transactions.len()];
40
41        // Index on transaction ids
42        let mut index = HashMap::with_capacity(transactions.len());
43        transactions.iter().enumerate().for_each(|(idx, tx)| {
44            let _ = index.insert(tx.as_ref().id(), idx);
45        });
46
47        // Transaction edges
48        let mut edges: Vec<Option<IndexSet>> = vec![None; transactions.len()];
49        transactions.iter().enumerate().for_each(|(destination_idx, tx)| {
50            tx.as_ref().inputs.iter().for_each(|input| {
51                if let Some(origin_idx) = index.get(&input.previous_outpoint.transaction_id) {
52                    edges[*origin_idx].get_or_insert_with(IndexSet::new).insert(destination_idx);
53                }
54            })
55        });
56
57        // Degrees
58        (0..transactions.len()).for_each(|origin_idx| {
59            if let Some(ref edges) = edges[origin_idx] {
60                edges.iter().for_each(|destination_idx| {
61                    #[allow(clippy::arithmetic_side_effects, reason = "ARITH-SAFETY(COUNTER)")]
62                    {
63                        in_degree[*destination_idx] += 1;
64                    }
65                });
66            }
67        });
68
69        // Degree 0
70        let mut queue = VecDeque::with_capacity(transactions.len());
71        (0..transactions.len()).for_each(|destination_idx| {
72            if in_degree[destination_idx] == 0 {
73                queue.push_back(destination_idx);
74            }
75        });
76        Self { transactions, in_degree, edges, queue, yields_count: 0 }
77    }
78}
79
80impl<'a, T: AsRef<Transaction>> Iterator for TopologicalIter<'a, T> {
81    type Item = &'a T;
82
83    fn next(&mut self) -> Option<Self::Item> {
84        match self.queue.pop_front() {
85            Some(current) => {
86                if let Some(ref edges) = self.edges[current] {
87                    edges.iter().for_each(|destination_idx| {
88                        let degree = self.in_degree.get_mut(*destination_idx).unwrap();
89                        #[allow(
90                            clippy::arithmetic_side_effects,
91                            reason = "Since the edge goes into `destination_idx`, the degree must be at least 1."
92                        )]
93                        {
94                            *degree -= 1;
95                        }
96                        if *degree == 0 {
97                            self.queue.push_back(*destination_idx);
98                        }
99                    });
100                }
101                #[allow(clippy::arithmetic_side_effects, reason = "ARITH-SAFETY(COUNTER)")]
102                {
103                    self.yields_count += 1;
104                }
105                Some(&self.transactions[current])
106            }
107            None => None,
108        }
109    }
110
111    fn size_hint(&self) -> (usize, Option<usize>) {
112        #[allow(
113            clippy::arithmetic_side_effects,
114            reason = "the subtraction is safe since we subtract at most `self.transactions.len()`."
115        )]
116        let items_remaining = self.transactions.len() - self.yields_count.min(self.transactions.len());
117        (self.yields_count, Some(items_remaining))
118    }
119}
120
121impl<T: AsRef<Transaction>> FusedIterator for TopologicalIter<'_, T> {}
122impl<T: AsRef<Transaction>> ExactSizeIterator for TopologicalIter<'_, T> {
123    fn len(&self) -> usize {
124        self.transactions.len()
125    }
126}
127
128pub trait IntoIterTopologically<T>
129where
130    T: AsRef<Transaction>,
131{
132    fn topological_into_iter(self) -> TopologicalIntoIter<T>;
133}
134
135impl<T: AsRef<Transaction>> IntoIterTopologically<T> for Vec<T> {
136    fn topological_into_iter(self) -> TopologicalIntoIter<T> {
137        TopologicalIntoIter::new(self)
138    }
139}
140
141impl<T, I, F> IntoIterTopologically<T> for Map<I, F>
142where
143    T: AsRef<Transaction>,
144    I: Iterator,
145    F: FnMut(<I as Iterator>::Item) -> T,
146{
147    fn topological_into_iter(self) -> TopologicalIntoIter<T> {
148        TopologicalIntoIter::new(self)
149    }
150}
151
152pub struct TopologicalIntoIter<T: AsRef<Transaction>> {
153    transactions: Vec<Option<T>>,
154    in_degree: Vec<isize>,
155    edges: Vec<Option<IndexSet>>,
156    queue: VecDeque<usize>,
157    yields_count: usize,
158}
159
160impl<T: AsRef<Transaction>> TopologicalIntoIter<T> {
161    pub fn new(transactions: impl IntoIterator<Item = T>) -> Self {
162        // Collect all transactions
163        let transactions = transactions.into_iter().map(|tx| Some(tx)).collect_vec();
164
165        let mut in_degree: Vec<isize> = vec![0; transactions.len()];
166
167        // Index on transaction ids
168        let mut index = HashMap::with_capacity(transactions.len());
169        transactions.iter().enumerate().for_each(|(idx, tx)| {
170            let _ = index.insert(tx.as_ref().unwrap().as_ref().id(), idx);
171        });
172
173        // Transaction edges
174        let mut edges: Vec<Option<IndexSet>> = vec![None; transactions.len()];
175        transactions.iter().enumerate().for_each(|(destination_idx, tx)| {
176            tx.as_ref().unwrap().as_ref().inputs.iter().for_each(|input| {
177                if let Some(origin_idx) = index.get(&input.previous_outpoint.transaction_id) {
178                    edges[*origin_idx].get_or_insert_with(IndexSet::new).insert(destination_idx);
179                }
180            })
181        });
182
183        // Degrees
184        (0..transactions.len()).for_each(|origin_idx| {
185            if let Some(ref edges) = edges[origin_idx] {
186                edges.iter().for_each(|destination_idx| {
187                    #[allow(clippy::arithmetic_side_effects, reason = "ARITH-SAFETY(COUNTER)")]
188                    {
189                        in_degree[*destination_idx] += 1;
190                    }
191                });
192            }
193        });
194
195        // Degree 0
196        let mut queue = VecDeque::with_capacity(transactions.len());
197        (0..transactions.len()).for_each(|destination_idx| {
198            if in_degree[destination_idx] == 0 {
199                queue.push_back(destination_idx);
200            }
201        });
202        Self { transactions, in_degree, edges, queue, yields_count: 0 }
203    }
204}
205
206impl<T: AsRef<Transaction>> Iterator for TopologicalIntoIter<T> {
207    type Item = T;
208
209    fn next(&mut self) -> Option<Self::Item> {
210        match self.queue.pop_front() {
211            Some(current) => {
212                if let Some(ref edges) = self.edges[current] {
213                    edges.iter().for_each(|destination_idx| {
214                        let degree = self.in_degree.get_mut(*destination_idx).unwrap();
215                        #[allow(
216                            clippy::arithmetic_side_effects,
217                            reason = "Since the edge goes into `destination_idx`, the degree must be at least 1."
218                        )]
219                        {
220                            *degree -= 1;
221                        }
222                        if *degree == 0 {
223                            self.queue.push_back(*destination_idx);
224                        }
225                    });
226                }
227                #[allow(clippy::arithmetic_side_effects, reason = "ARITH-SAFETY(COUNTER)")]
228                {
229                    self.yields_count += 1;
230                }
231
232                self.transactions[current].take()
233            }
234            None => None,
235        }
236    }
237
238    fn size_hint(&self) -> (usize, Option<usize>) {
239        #[allow(
240            clippy::arithmetic_side_effects,
241            reason = "the subtractin is safe since we subtract at most `self.transactions.len()`."
242        )]
243        let items_remaining = self.transactions.len() - self.yields_count.min(self.transactions.len());
244        (self.yields_count, Some(items_remaining))
245    }
246}
247
248impl<T: AsRef<Transaction>> FusedIterator for TopologicalIntoIter<T> {}
249impl<T: AsRef<Transaction>> ExactSizeIterator for TopologicalIntoIter<T> {
250    fn len(&self) -> usize {
251        self.transactions.len()
252    }
253}