Preprocess dominator tree to answer queries in O(1)

7 files changed, 137 insertions, 55 deletions

diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
index 3c350e25ba6..2d8cbea8bac 100644
--- a/compiler/rustc_const_eval/src/transform/validate.rs
+++ b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -155,7 +155,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
         if self.unwind_edge_count <= 1 {
             return;
         }
-        let doms = self.body.basic_blocks.dominators();
+        let dom_tree = self.body.basic_blocks.dominator_tree();
         let mut post_contract_node = FxHashMap::default();
         // Reusing the allocation across invocations of the closure
         let mut dom_path = vec![];
@@ -164,7 +164,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 if let Some(root) = post_contract_node.get(&bb) {
                     break *root;
                 }
-                let parent = doms.immediate_dominator(bb).unwrap();
+                let parent = dom_tree.immediate_dominator(bb).unwrap();
                 dom_path.push(bb);
                 if !self.body.basic_blocks[parent].is_cleanup {
                     break bb;
diff --git a/compiler/rustc_data_structures/src/graph/dominators/mod.rs b/compiler/rustc_data_structures/src/graph/dominators/mod.rs
index a7de709ba72..5d6a1de1d21 100644
--- a/compiler/rustc_data_structures/src/graph/dominators/mod.rs
+++ b/compiler/rustc_data_structures/src/graph/dominators/mod.rs
@@ -26,7 +26,7 @@ rustc_index::newtype_index! {
     struct PreorderIndex {}
 }
 
-pub fn dominators<G: ControlFlowGraph>(graph: G) -> Dominators<G::Node> {
+pub fn dominator_tree<G: ControlFlowGraph>(graph: G) -> DominatorTree<G::Node> {
     // compute the post order index (rank) for each node
     let mut post_order_rank = IndexVec::from_elem_n(0, graph.num_nodes());
 
@@ -244,7 +244,7 @@ pub fn dominators<G: ControlFlowGraph>(graph: G) -> Dominators<G::Node> {
 
     let start_node = graph.start_node();
     immediate_dominators[start_node] = None;
-    Dominators { start_node, post_order_rank, immediate_dominators }
+    DominatorTree { start_node, post_order_rank, immediate_dominators }
 }
 
 /// Evaluate the link-eval virtual forest, providing the currently minimum semi
@@ -309,16 +309,18 @@ fn compress(
 
 /// Tracks the list of dominators for each node.
 #[derive(Clone, Debug)]
-pub struct Dominators<N: Idx> {
+pub struct DominatorTree<N: Idx> {
     start_node: N,
     post_order_rank: IndexVec<N, usize>,
     // Even though we track only the immediate dominator of each node, it's
     // possible to get its full list of dominators by looking up the dominator
     // of each dominator. (See the `impl Iterator for Iter` definition).
+    //
+    // Note: immediate_dominators[root] is Some(root)!
     immediate_dominators: IndexVec<N, Option<N>>,
 }
 
-impl<Node: Idx> Dominators<Node> {
+impl<Node: Idx> DominatorTree<Node> {
     /// Returns true if node is reachable from the start node.
     pub fn is_reachable(&self, node: Node) -> bool {
         node == self.start_node || self.immediate_dominators[node].is_some()
@@ -333,12 +335,7 @@ impl<Node: Idx> Dominators<Node> {
     /// See the `impl Iterator for Iter` definition to understand how this works.
     pub fn dominators(&self, node: Node) -> Iter<'_, Node> {
         assert!(self.is_reachable(node), "node {node:?} is not reachable");
-        Iter { dominators: self, node: Some(node) }
-    }
-
-    pub fn dominates(&self, dom: Node, node: Node) -> bool {
-        // FIXME -- could be optimized by using post-order-rank
-        self.dominators(node).any(|n| n == dom)
+        Iter { dom_tree: self, node: Some(node) }
     }
 
     /// Provide deterministic ordering of nodes such that, if any two nodes have a dominator
@@ -351,7 +348,7 @@ impl<Node: Idx> Dominators<Node> {
 }
 
 pub struct Iter<'dom, Node: Idx> {
-    dominators: &'dom Dominators<Node>,
+    dom_tree: &'dom DominatorTree<Node>,
     node: Option<Node>,
 }
 
@@ -360,10 +357,96 @@ impl<'dom, Node: Idx> Iterator for Iter<'dom, Node> {
 
     fn next(&mut self) -> Option<Self::Item> {
         if let Some(node) = self.node {
-            self.node = self.dominators.immediate_dominator(node);
+            self.node = self.dom_tree.immediate_dominator(node);
             Some(node)
         } else {
             None
         }
     }
 }
+
+#[derive(Clone, Debug)]
+pub struct Dominators<Node: Idx> {
+    time: IndexVec<Node, Time>,
+}
+
+/// Describes the number of vertices discovered at the time when processing of a particular vertex
+/// started and when it finished. Both values are zero for unreachable vertices.
+#[derive(Copy, Clone, Default, Debug)]
+struct Time {
+    start: u32,
+    finish: u32,
+}
+
+impl<Node: Idx> Dominators<Node> {
+    pub fn dummy() -> Self {
+        Self { time: Default::default() }
+    }
+
+    /// Returns true if `a` dominates `b`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `b` is unreachable.
+    pub fn dominates(&self, a: Node, b: Node) -> bool {
+        let a = self.time[a];
+        let b = self.time[b];
+        assert!(b.start != 0, "node {b:?} is not reachable");
+        a.start <= b.start && b.finish <= a.finish
+    }
+}
+
+pub fn dominators<N: Idx>(tree: &DominatorTree<N>) -> Dominators<N> {
+    let DominatorTree { start_node, ref immediate_dominators, post_order_rank: _ } = *tree;
+
+    // Transpose the dominator tree edges, so that child nodes of vertex v are stored in
+    // node[edges[v].start..edges[y].end].
+    let mut edges: IndexVec<N, std::ops::Range<u32>> =
+        IndexVec::from_elem(0..0, immediate_dominators);
+    for &idom in immediate_dominators.iter() {
+        if let Some(idom) = idom {
+            edges[idom].end += 1;
+        }
+    }
+    let mut m = 0;
+    for e in edges.iter_mut() {
+        m += e.end;
+        e.start = m;
+        e.end = m;
+    }
+    let mut node = IndexVec::from_elem_n(Idx::new(0), m.try_into().unwrap());
+    for (i, &idom) in immediate_dominators.iter_enumerated() {
+        if let Some(idom) = idom {
+            edges[idom].start -= 1;
+            node[edges[idom].start] = i;
+        }
+    }
+
+    // Perform a depth-first search of the dominator tree. Record the number of vertices discovered
+    // when vertex v is discovered first as time[v].start, and when its processing is finished as
+    // time[v].finish.
+    let mut time: IndexVec<N, Time> = IndexVec::from_elem(Time::default(), immediate_dominators);
+    let mut stack = Vec::new();
+
+    let mut discovered = 1;
+    stack.push(start_node);
+    time[start_node].start = discovered;
+
+    while let Some(&i) = stack.last() {
+        let e = &mut edges[i];
+        if e.start == e.end {
+            // Finish processing vertex i.
+            time[i].finish = discovered;
+            stack.pop();
+        } else {
+            let j = node[e.start];
+            e.start += 1;
+            // Start processing vertex j.
+            discovered += 1;
+            time[j].start = discovered;
+            stack.push(j);
+        }
+    }
+
+    Dominators { time }
+}
diff --git a/compiler/rustc_data_structures/src/graph/dominators/tests.rs b/compiler/rustc_data_structures/src/graph/dominators/tests.rs
index 8b124516623..61a21724dda 100644
--- a/compiler/rustc_data_structures/src/graph/dominators/tests.rs
+++ b/compiler/rustc_data_structures/src/graph/dominators/tests.rs
@@ -6,8 +6,8 @@ use super::super::tests::TestGraph;
 fn diamond() {
     let graph = TestGraph::new(0, &[(0, 1), (0, 2), (1, 3), (2, 3)]);
 
-    let dominators = dominators(&graph);
-    let immediate_dominators = &dominators.immediate_dominators;
+    let tree = dominator_tree(&graph);
+    let immediate_dominators = &tree.immediate_dominators;
     assert_eq!(immediate_dominators[0], None);
     assert_eq!(immediate_dominators[1], Some(0));
     assert_eq!(immediate_dominators[2], Some(0));
@@ -22,8 +22,8 @@ fn paper() {
         &[(6, 5), (6, 4), (5, 1), (4, 2), (4, 3), (1, 2), (2, 3), (3, 2), (2, 1)],
     );
 
-    let dominators = dominators(&graph);
-    let immediate_dominators = &dominators.immediate_dominators;
+    let dom_tree = dominator_tree(&graph);
+    let immediate_dominators = &dom_tree.immediate_dominators;
     assert_eq!(immediate_dominators[0], None); // <-- note that 0 is not in graph
     assert_eq!(immediate_dominators[1], Some(6));
     assert_eq!(immediate_dominators[2], Some(6));
@@ -41,15 +41,15 @@ fn paper_slt() {
         &[(1, 2), (1, 3), (2, 3), (2, 7), (3, 4), (3, 6), (4, 5), (5, 4), (6, 7), (7, 8), (8, 5)],
     );
 
-    dominators(&graph);
+    dominator_tree(&graph);
 }
 
 #[test]
 fn immediate_dominator() {
     let graph = TestGraph::new(1, &[(1, 2), (2, 3)]);
-    let dominators = dominators(&graph);
-    assert_eq!(dominators.immediate_dominator(0), None);
-    assert_eq!(dominators.immediate_dominator(1), None);
-    assert_eq!(dominators.immediate_dominator(2), Some(1));
-    assert_eq!(dominators.immediate_dominator(3), Some(2));
+    let tree = dominator_tree(&graph);
+    assert_eq!(tree.immediate_dominator(0), None);
+    assert_eq!(tree.immediate_dominator(1), None);
+    assert_eq!(tree.immediate_dominator(2), Some(1));
+    assert_eq!(tree.immediate_dominator(3), Some(2));
 }
diff --git a/compiler/rustc_middle/src/mir/basic_blocks.rs b/compiler/rustc_middle/src/mir/basic_blocks.rs
index 1319ddbb877..f5157e18e6d 100644
--- a/compiler/rustc_middle/src/mir/basic_blocks.rs
+++ b/compiler/rustc_middle/src/mir/basic_blocks.rs
@@ -3,6 +3,7 @@ use crate::mir::{BasicBlock, BasicBlockData, Successors, Terminator, TerminatorK
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::graph;
+use rustc_data_structures::graph::dominators::{dominator_tree, DominatorTree};
 use rustc_data_structures::graph::dominators::{dominators, Dominators};
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::OnceCell;
@@ -41,8 +42,12 @@ impl<'tcx> BasicBlocks<'tcx> {
         *self.cache.is_cyclic.get_or_init(|| graph::is_cyclic(self))
     }
 
+    pub fn dominator_tree(&self) -> DominatorTree<BasicBlock> {
+        dominator_tree(&self)
+    }
+
     pub fn dominators(&self) -> Dominators<BasicBlock> {
-        dominators(&self)
+        dominators(&self.dominator_tree())
     }
 
     /// Returns predecessors for each basic block.
diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs
index 986d2fd190d..0126310e9ff 100644
--- a/compiler/rustc_mir_transform/src/coverage/graph.rs
+++ b/compiler/rustc_mir_transform/src/coverage/graph.rs
@@ -2,13 +2,14 @@ use super::Error;
 
 use itertools::Itertools;
 use rustc_data_structures::fx::FxHashMap;
-use rustc_data_structures::graph::dominators::{self, Dominators};
+use rustc_data_structures::graph::dominators::{self, DominatorTree, Dominators};
 use rustc_data_structures::graph::{self, GraphSuccessors, WithNumNodes, WithStartNode};
 use rustc_index::bit_set::BitSet;
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::mir::coverage::*;
 use rustc_middle::mir::{self, BasicBlock, BasicBlockData, Terminator, TerminatorKind};
 
+use std::cmp::Ordering;
 use std::ops::{Index, IndexMut};
 
 const ID_SEPARATOR: &str = ",";
@@ -24,6 +25,7 @@ pub(super) struct CoverageGraph {
     bb_to_bcb: IndexVec<BasicBlock, Option<BasicCoverageBlock>>,
     pub successors: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,
     pub predecessors: IndexVec<BasicCoverageBlock, Vec<BasicCoverageBlock>>,
+    dominator_tree: Option<DominatorTree<BasicCoverageBlock>>,
     dominators: Option<Dominators<BasicCoverageBlock>>,
 }
 
@@ -66,9 +68,17 @@ impl CoverageGraph {
             }
         }
 
-        let mut basic_coverage_blocks =
-            Self { bcbs, bb_to_bcb, successors, predecessors, dominators: None };
-        let dominators = dominators::dominators(&basic_coverage_blocks);
+        let mut basic_coverage_blocks = Self {
+            bcbs,
+            bb_to_bcb,
+            successors,
+            predecessors,
+            dominator_tree: None,
+            dominators: None,
+        };
+        let dominator_tree = dominators::dominator_tree(&basic_coverage_blocks);
+        let dominators = dominators::dominators(&dominator_tree);
+        basic_coverage_blocks.dominator_tree = Some(dominator_tree);
         basic_coverage_blocks.dominators = Some(dominators);
         basic_coverage_blocks
     }
@@ -212,8 +222,12 @@ impl CoverageGraph {
     }
 
     #[inline(always)]
-    pub fn dominators(&self) -> &Dominators<BasicCoverageBlock> {
-        self.dominators.as_ref().unwrap()
+    pub fn rank_partial_cmp(
+        &self,
+        a: BasicCoverageBlock,
+        b: BasicCoverageBlock,
+    ) -> Option<Ordering> {
+        self.dominator_tree.as_ref().unwrap().rank_partial_cmp(a, b)
     }
 }
 
@@ -650,26 +664,6 @@ pub(super) fn find_loop_backedges(
     let mut backedges = IndexVec::from_elem_n(Vec::<BasicCoverageBlock>::new(), num_bcbs);
 
     // Identify loops by their backedges.
-    //
-    // The computational complexity is bounded by: n(s) x d where `n` is the number of
-    // `BasicCoverageBlock` nodes (the simplified/reduced representation of the CFG derived from the
-    // MIR); `s` is the average number of successors per node (which is most likely less than 2, and
-    // independent of the size of the function, so it can be treated as a constant);
-    // and `d` is the average number of dominators per node.
-    //
-    // The average number of dominators depends on the size and complexity of the function, and
-    // nodes near the start of the function's control flow graph typically have less dominators
-    // than nodes near the end of the CFG. Without doing a detailed mathematical analysis, I
-    // think the resulting complexity has the characteristics of O(n log n).
-    //
-    // The overall complexity appears to be comparable to many other MIR transform algorithms, and I
-    // don't expect that this function is creating a performance hot spot, but if this becomes an
-    // issue, there may be ways to optimize the `dominates` algorithm (as indicated by an
-    // existing `FIXME` comment in that code), or possibly ways to optimize it's usage here, perhaps
-    // by keeping track of results for visited `BasicCoverageBlock`s if they can be used to short
-    // circuit downstream `dominates` checks.
-    //
-    // For now, that kind of optimization seems unnecessarily complicated.
     for (bcb, _) in basic_coverage_blocks.iter_enumerated() {
         for &successor in &basic_coverage_blocks.successors[bcb] {
             if basic_coverage_blocks.dominates(successor, bcb) {
diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs
index 287ae217087..97dd1dd09a2 100644
--- a/compiler/rustc_mir_transform/src/coverage/spans.rs
+++ b/compiler/rustc_mir_transform/src/coverage/spans.rs
@@ -345,7 +345,7 @@ impl<'a, 'tcx> CoverageSpans<'a, 'tcx> {
                         // before the dominated equal spans). When later comparing two spans in
                         // order, the first will either dominate the second, or they will have no
                         // dominator relationship.
-                        self.basic_coverage_blocks.dominators().rank_partial_cmp(a.bcb, b.bcb)
+                        self.basic_coverage_blocks.rank_partial_cmp(a.bcb, b.bcb)
                     }
                 } else {
                     // Sort hi() in reverse order so shorter spans are attempted after longer spans.
diff --git a/compiler/rustc_mir_transform/src/ctfe_limit.rs b/compiler/rustc_mir_transform/src/ctfe_limit.rs
index 1b3ac78fbc6..a5b3873533c 100644
--- a/compiler/rustc_mir_transform/src/ctfe_limit.rs
+++ b/compiler/rustc_mir_transform/src/ctfe_limit.rs
@@ -2,7 +2,7 @@
 //! (thus indicating there is a loop in the CFG), or whose terminator is a function call.
 use crate::MirPass;
 
-use rustc_data_structures::graph::dominators::Dominators;
+use rustc_data_structures::graph::dominators::DominatorTree;
 use rustc_middle::mir::{
     BasicBlock, BasicBlockData, Body, Statement, StatementKind, TerminatorKind,
 };
@@ -13,7 +13,7 @@ pub struct CtfeLimit;
 impl<'tcx> MirPass<'tcx> for CtfeLimit {
     #[instrument(skip(self, _tcx, body))]
     fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        let doms = body.basic_blocks.dominators();
+        let doms = body.basic_blocks.dominator_tree();
         let indices: Vec<BasicBlock> = body
             .basic_blocks
             .iter_enumerated()
@@ -39,7 +39,7 @@ impl<'tcx> MirPass<'tcx> for CtfeLimit {
 }
 
 fn has_back_edge(
-    doms: &Dominators<BasicBlock>,
+    doms: &DominatorTree<BasicBlock>,
     node: BasicBlock,
     node_data: &BasicBlockData<'_>,
 ) -> bool {