1use alloc::{vec, vec::Vec};
16use core::{cmp::Ordering, hash::Hash};
17
18use hashbrown::{hash_map::Iter, HashMap, HashSet};
19
20use crate::visit::{DfsPostOrder, GraphBase, IntoNeighbors, Visitable, Walker};
21
22#[derive(Debug, Clone)]
24pub struct Dominators<N>
25where
26 N: Copy + Eq + Hash,
27{
28 root: N,
29 dominators: HashMap<N, N>,
30}
31
32impl<N> Dominators<N>
33where
34 N: Copy + Eq + Hash,
35{
36 pub fn root(&self) -> N {
38 self.root
39 }
40
41 pub fn immediate_dominator(&self, node: N) -> Option<N> {
46 if node == self.root {
47 None
48 } else {
49 self.dominators.get(&node).cloned()
50 }
51 }
52
53 pub fn strict_dominators(&self, node: N) -> Option<DominatorsIter<N>> {
58 if self.dominators.contains_key(&node) {
59 Some(DominatorsIter {
60 dominators: self,
61 node: self.immediate_dominator(node),
62 })
63 } else {
64 None
65 }
66 }
67
68 pub fn dominators(&self, node: N) -> Option<DominatorsIter<N>> {
74 if self.dominators.contains_key(&node) {
75 Some(DominatorsIter {
76 dominators: self,
77 node: Some(node),
78 })
79 } else {
80 None
81 }
82 }
83
84 pub fn immediately_dominated_by(&self, node: N) -> DominatedByIter<N> {
87 DominatedByIter {
88 iter: self.dominators.iter(),
89 node,
90 }
91 }
92}
93
94#[derive(Debug, Clone)]
96pub struct DominatorsIter<'a, N>
97where
98 N: 'a + Copy + Eq + Hash,
99{
100 dominators: &'a Dominators<N>,
101 node: Option<N>,
102}
103
104impl<'a, N> Iterator for DominatorsIter<'a, N>
105where
106 N: 'a + Copy + Eq + Hash,
107{
108 type Item = N;
109
110 fn next(&mut self) -> Option<Self::Item> {
111 let next = self.node.take();
112 if let Some(next) = next {
113 self.node = self.dominators.immediate_dominator(next);
114 }
115 next
116 }
117}
118
119#[derive(Debug, Clone)]
121pub struct DominatedByIter<'a, N>
122where
123 N: 'a + Copy + Eq + Hash,
124{
125 iter: Iter<'a, N, N>,
126 node: N,
127}
128
129impl<'a, N> Iterator for DominatedByIter<'a, N>
130where
131 N: 'a + Copy + Eq + Hash,
132{
133 type Item = N;
134
135 fn next(&mut self) -> Option<Self::Item> {
136 for (dominator, dominated) in self.iter.by_ref() {
137 if dominated == &self.node && dominated != dominator {
140 return Some(*dominator);
141 }
142 }
143 None
144 }
145 fn size_hint(&self) -> (usize, Option<usize>) {
146 let (_, upper) = self.iter.size_hint();
147 (0, upper)
148 }
149}
150
151const UNDEFINED: usize = usize::MAX;
154
155pub fn simple_fast<G>(graph: G, root: G::NodeId) -> Dominators<G::NodeId>
165where
166 G: IntoNeighbors + Visitable,
167 <G as GraphBase>::NodeId: Eq + Hash,
168{
169 let (post_order, predecessor_sets) = simple_fast_post_order(graph, root);
170 let length = post_order.len();
171 debug_assert!(length > 0);
172 debug_assert!(post_order.last() == Some(&root));
173
174 let node_to_post_order_idx: HashMap<_, _> = post_order
181 .iter()
182 .enumerate()
183 .map(|(idx, &node)| (node, idx))
184 .collect();
185
186 let idx_to_predecessor_vec =
189 predecessor_sets_to_idx_vecs(&post_order, &node_to_post_order_idx, predecessor_sets);
190
191 let mut dominators = vec![UNDEFINED; length];
192 dominators[length - 1] = length - 1;
193
194 let mut changed = true;
195 while changed {
196 changed = false;
197
198 for idx in (0..length - 1).rev() {
201 debug_assert!(post_order[idx] != root);
202
203 let new_idom_idx = {
208 let mut predecessors = idx_to_predecessor_vec[idx]
209 .iter()
210 .filter(|&&p| dominators[p] != UNDEFINED);
211 let new_idom_idx = predecessors.next().expect(
212 "Because the root is initialized to dominate itself, and is the \
213 first node in every path, there must exist a predecessor to this \
214 node that also has a dominator",
215 );
216 predecessors.fold(*new_idom_idx, |new_idom_idx, &predecessor_idx| {
217 intersect(&dominators, new_idom_idx, predecessor_idx)
218 })
219 };
220
221 debug_assert!(new_idom_idx < length);
222
223 if new_idom_idx != dominators[idx] {
224 dominators[idx] = new_idom_idx;
225 changed = true;
226 }
227 }
228 }
229
230 debug_assert!(!dominators.iter().any(|&dom| dom == UNDEFINED));
233
234 Dominators {
235 root,
236 dominators: dominators
237 .into_iter()
238 .enumerate()
239 .map(|(idx, dom_idx)| (post_order[idx], post_order[dom_idx]))
240 .collect(),
241 }
242}
243
244fn intersect(dominators: &[usize], mut finger1: usize, mut finger2: usize) -> usize {
245 loop {
246 match finger1.cmp(&finger2) {
247 Ordering::Less => finger1 = dominators[finger1],
248 Ordering::Greater => finger2 = dominators[finger2],
249 Ordering::Equal => return finger1,
250 }
251 }
252}
253
254fn predecessor_sets_to_idx_vecs<N>(
255 post_order: &[N],
256 node_to_post_order_idx: &HashMap<N, usize>,
257 mut predecessor_sets: HashMap<N, HashSet<N>>,
258) -> Vec<Vec<usize>>
259where
260 N: Copy + Eq + Hash,
261{
262 post_order
263 .iter()
264 .map(|node| {
265 predecessor_sets
266 .remove(node)
267 .map(|predecessors| {
268 predecessors
269 .into_iter()
270 .map(|p| *node_to_post_order_idx.get(&p).unwrap())
271 .collect()
272 })
273 .unwrap_or_default()
274 })
275 .collect()
276}
277
278type PredecessorSets<NodeId> = HashMap<NodeId, HashSet<NodeId>>;
279
280fn simple_fast_post_order<G>(
281 graph: G,
282 root: G::NodeId,
283) -> (Vec<G::NodeId>, PredecessorSets<G::NodeId>)
284where
285 G: IntoNeighbors + Visitable,
286 <G as GraphBase>::NodeId: Eq + Hash,
287{
288 let mut post_order = vec![];
289 let mut predecessor_sets = HashMap::new();
290
291 for node in DfsPostOrder::new(graph, root).iter(graph) {
292 post_order.push(node);
293
294 for successor in graph.neighbors(node) {
295 predecessor_sets
296 .entry(successor)
297 .or_insert_with(HashSet::new)
298 .insert(node);
299 }
300 }
301
302 (post_order, predecessor_sets)
303}
304
305#[cfg(test)]
306mod tests {
307 use super::*;
308
309 #[test]
310 fn test_iter_dominators() {
311 let doms: Dominators<u32> = Dominators {
312 root: 0,
313 dominators: [(2, 1), (1, 0), (0, 0)].iter().cloned().collect(),
314 };
315
316 let all_doms: Vec<_> = doms.dominators(2).unwrap().collect();
317 assert_eq!(vec![2, 1, 0], all_doms);
318
319 assert_eq!(None::<()>, doms.dominators(99).map(|_| unreachable!()));
320
321 let strict_doms: Vec<_> = doms.strict_dominators(2).unwrap().collect();
322 assert_eq!(vec![1, 0], strict_doms);
323
324 assert_eq!(
325 None::<()>,
326 doms.strict_dominators(99).map(|_| unreachable!())
327 );
328
329 let dom_by: Vec<_> = doms.immediately_dominated_by(1).collect();
330 assert_eq!(vec![2], dom_by);
331 assert_eq!(None, doms.immediately_dominated_by(99).next());
332 assert_eq!(1, doms.immediately_dominated_by(0).count());
333 }
334}