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>
179where
180 G: IntoNeighbors + Visitable,
181 <G as GraphBase>::NodeId: Eq + Hash,
182{
183 let (post_order, predecessor_sets) = simple_fast_post_order(graph, root);
184 let length = post_order.len();
185 debug_assert!(length > 0);
186 debug_assert!(post_order.last() == Some(&root));
187
188 let node_to_post_order_idx: HashMap<_, _> = post_order
195 .iter()
196 .enumerate()
197 .map(|(idx, &node)| (node, idx))
198 .collect();
199
200 let idx_to_predecessor_vec =
203 predecessor_sets_to_idx_vecs(&post_order, &node_to_post_order_idx, predecessor_sets);
204
205 let mut dominators = vec![UNDEFINED; length];
206 dominators[length - 1] = length - 1;
207
208 let mut changed = true;
209 while changed {
210 changed = false;
211
212 for idx in (0..length - 1).rev() {
215 debug_assert!(post_order[idx] != root);
216
217 let new_idom_idx = {
222 let mut predecessors = idx_to_predecessor_vec[idx]
223 .iter()
224 .filter(|&&p| dominators[p] != UNDEFINED);
225 let new_idom_idx = predecessors.next().expect(
226 "Because the root is initialized to dominate itself, and is the \
227 first node in every path, there must exist a predecessor to this \
228 node that also has a dominator",
229 );
230 predecessors.fold(*new_idom_idx, |new_idom_idx, &predecessor_idx| {
231 intersect(&dominators, new_idom_idx, predecessor_idx)
232 })
233 };
234
235 debug_assert!(new_idom_idx < length);
236
237 if new_idom_idx != dominators[idx] {
238 dominators[idx] = new_idom_idx;
239 changed = true;
240 }
241 }
242 }
243
244 debug_assert!(!dominators.contains(&UNDEFINED));
246
247 Dominators {
248 root,
249 dominators: dominators
250 .into_iter()
251 .enumerate()
252 .map(|(idx, dom_idx)| (post_order[idx], post_order[dom_idx]))
253 .collect(),
254 }
255}
256
257fn intersect(dominators: &[usize], mut finger1: usize, mut finger2: usize) -> usize {
258 loop {
259 match finger1.cmp(&finger2) {
260 Ordering::Less => finger1 = dominators[finger1],
261 Ordering::Greater => finger2 = dominators[finger2],
262 Ordering::Equal => return finger1,
263 }
264 }
265}
266
267fn predecessor_sets_to_idx_vecs<N>(
268 post_order: &[N],
269 node_to_post_order_idx: &HashMap<N, usize>,
270 mut predecessor_sets: HashMap<N, HashSet<N>>,
271) -> Vec<Vec<usize>>
272where
273 N: Copy + Eq + Hash,
274{
275 post_order
276 .iter()
277 .map(|node| {
278 predecessor_sets
279 .remove(node)
280 .map(|predecessors| {
281 predecessors
282 .into_iter()
283 .map(|p| *node_to_post_order_idx.get(&p).unwrap())
284 .collect()
285 })
286 .unwrap_or_default()
287 })
288 .collect()
289}
290
291type PredecessorSets<NodeId> = HashMap<NodeId, HashSet<NodeId>>;
292
293fn simple_fast_post_order<G>(
294 graph: G,
295 root: G::NodeId,
296) -> (Vec<G::NodeId>, PredecessorSets<G::NodeId>)
297where
298 G: IntoNeighbors + Visitable,
299 <G as GraphBase>::NodeId: Eq + Hash,
300{
301 let mut post_order = vec![];
302 let mut predecessor_sets = HashMap::new();
303
304 for node in DfsPostOrder::new(graph, root).iter(graph) {
305 post_order.push(node);
306
307 for successor in graph.neighbors(node) {
308 predecessor_sets
309 .entry(successor)
310 .or_insert_with(HashSet::new)
311 .insert(node);
312 }
313 }
314
315 (post_order, predecessor_sets)
316}
317
318#[cfg(test)]
319mod tests {
320 use super::*;
321
322 #[test]
323 fn test_iter_dominators() {
324 let doms: Dominators<u32> = Dominators {
325 root: 0,
326 dominators: [(2, 1), (1, 0), (0, 0)].iter().cloned().collect(),
327 };
328
329 let all_doms: Vec<_> = doms.dominators(2).unwrap().collect();
330 assert_eq!(vec![2, 1, 0], all_doms);
331
332 assert_eq!(None::<()>, doms.dominators(99).map(|_| unreachable!()));
333
334 let strict_doms: Vec<_> = doms.strict_dominators(2).unwrap().collect();
335 assert_eq!(vec![1, 0], strict_doms);
336
337 assert_eq!(
338 None::<()>,
339 doms.strict_dominators(99).map(|_| unreachable!())
340 );
341
342 let dom_by: Vec<_> = doms.immediately_dominated_by(1).collect();
343 assert_eq!(vec![2], dom_by);
344 assert_eq!(None, doms.immediately_dominated_by(99).next());
345 assert_eq!(1, doms.immediately_dominated_by(0).count());
346 }
347}