Auto merge of #116454 - tmiasko:small-dominators, r=cjgillot

Generalize small dominators optimization

* Use small dominators optimization from 640ede7b0a1840415cb6ec881c2210302bfeba18 more generally.
* Merge `DefLocation` and `LocationExtended` since they serve the same purpose.

5 files changed, 134 insertions, 115 deletions

diff --git a/compiler/rustc_codegen_ssa/src/mir/analyze.rs b/compiler/rustc_codegen_ssa/src/mir/analyze.rs
index 7fda2d5fadf..62e997e5cfa 100644
--- a/compiler/rustc_codegen_ssa/src/mir/analyze.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/analyze.rs
@@ -8,7 +8,7 @@ use rustc_index::bit_set::BitSet;
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::mir::traversal;
 use rustc_middle::mir::visit::{MutatingUseContext, NonMutatingUseContext, PlaceContext, Visitor};
-use rustc_middle::mir::{self, Location, TerminatorKind};
+use rustc_middle::mir::{self, DefLocation, Location, TerminatorKind};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf};
 
 pub fn non_ssa_locals<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
@@ -67,21 +67,6 @@ enum LocalKind {
     SSA(DefLocation),
 }
 
-#[derive(Copy, Clone, PartialEq, Eq)]
-enum DefLocation {
-    Argument,
-    Body(Location),
-}
-
-impl DefLocation {
-    fn dominates(self, location: Location, dominators: &Dominators<mir::BasicBlock>) -> bool {
-        match self {
-            DefLocation::Argument => true,
-            DefLocation::Body(def) => def.successor_within_block().dominates(location, dominators),
-        }
-    }
-}
-
 struct LocalAnalyzer<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> {
     fx: &'mir FunctionCx<'a, 'tcx, Bx>,
     dominators: &'mir Dominators<mir::BasicBlock>,
diff --git a/compiler/rustc_data_structures/src/graph/dominators/mod.rs b/compiler/rustc_data_structures/src/graph/dominators/mod.rs
index 4075481e561..9685ad24a97 100644
--- a/compiler/rustc_data_structures/src/graph/dominators/mod.rs
+++ b/compiler/rustc_data_structures/src/graph/dominators/mod.rs
@@ -26,7 +26,42 @@ rustc_index::newtype_index! {
     struct PreorderIndex {}
 }
 
-pub fn dominators<G: ControlFlowGraph>(graph: &G) -> Dominators<G::Node> {
+#[derive(Clone, Debug)]
+pub struct Dominators<Node: Idx> {
+    kind: Kind<Node>,
+}
+
+#[derive(Clone, Debug)]
+enum Kind<Node: Idx> {
+    /// A representation optimized for a small path graphs.
+    Path,
+    General(Inner<Node>),
+}
+
+pub fn dominators<G: ControlFlowGraph>(g: &G) -> Dominators<G::Node> {
+    // We often encounter MIR bodies with 1 or 2 basic blocks. Special case the dominators
+    // computation and representation for those cases.
+    if is_small_path_graph(g) {
+        Dominators { kind: Kind::Path }
+    } else {
+        Dominators { kind: Kind::General(dominators_impl(g)) }
+    }
+}
+
+fn is_small_path_graph<G: ControlFlowGraph>(g: &G) -> bool {
+    if g.start_node().index() != 0 {
+        return false;
+    }
+    if g.num_nodes() == 1 {
+        return true;
+    }
+    if g.num_nodes() == 2 {
+        return g.successors(g.start_node()).any(|n| n.index() == 1);
+    }
+    false
+}
+
+fn dominators_impl<G: ControlFlowGraph>(graph: &G) -> Inner<G::Node> {
     // compute the post order index (rank) for each node
     let mut post_order_rank = IndexVec::from_elem_n(0, graph.num_nodes());
 
@@ -245,7 +280,7 @@ pub fn dominators<G: ControlFlowGraph>(graph: &G) -> Dominators<G::Node> {
 
     let time = compute_access_time(start_node, &immediate_dominators);
 
-    Dominators { start_node, post_order_rank, immediate_dominators, time }
+    Inner { post_order_rank, immediate_dominators, time }
 }
 
 /// Evaluate the link-eval virtual forest, providing the currently minimum semi
@@ -310,8 +345,7 @@ fn compress(
 
 /// Tracks the list of dominators for each node.
 #[derive(Clone, Debug)]
-pub struct Dominators<N: Idx> {
-    start_node: N,
+struct Inner<N: Idx> {
     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
@@ -323,12 +357,24 @@ pub struct Dominators<N: Idx> {
 impl<Node: Idx> Dominators<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()
+        match &self.kind {
+            Kind::Path => true,
+            Kind::General(g) => g.time[node].start != 0,
+        }
     }
 
     /// Returns the immediate dominator of node, if any.
     pub fn immediate_dominator(&self, node: Node) -> Option<Node> {
-        self.immediate_dominators[node]
+        match &self.kind {
+            Kind::Path => {
+                if 0 < node.index() {
+                    Some(Node::new(node.index() - 1))
+                } else {
+                    None
+                }
+            }
+            Kind::General(g) => g.immediate_dominators[node],
+        }
     }
 
     /// Provides an iterator over each dominator up the CFG, for the given Node.
@@ -343,7 +389,10 @@ impl<Node: Idx> Dominators<Node> {
     /// of two unrelated nodes will also be consistent, but otherwise the order has no
     /// meaning.) This method cannot be used to determine if either Node dominates the other.
     pub fn cmp_in_dominator_order(&self, lhs: Node, rhs: Node) -> Ordering {
-        self.post_order_rank[rhs].cmp(&self.post_order_rank[lhs])
+        match &self.kind {
+            Kind::Path => lhs.index().cmp(&rhs.index()),
+            Kind::General(g) => g.post_order_rank[rhs].cmp(&g.post_order_rank[lhs]),
+        }
     }
 
     /// Returns true if `a` dominates `b`.
@@ -352,10 +401,15 @@ impl<Node: Idx> Dominators<Node> {
     ///
     /// 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
+        match &self.kind {
+            Kind::Path => a.index() <= b.index(),
+            Kind::General(g) => {
+                let a = g.time[a];
+                let b = g.time[b];
+                assert!(b.start != 0, "node {b:?} is not reachable");
+                a.start <= b.start && b.finish <= a.finish
+            }
+        }
     }
 }
 
diff --git a/compiler/rustc_data_structures/src/graph/dominators/tests.rs b/compiler/rustc_data_structures/src/graph/dominators/tests.rs
index 5472bb8087e..39725ba4301 100644
--- a/compiler/rustc_data_structures/src/graph/dominators/tests.rs
+++ b/compiler/rustc_data_structures/src/graph/dominators/tests.rs
@@ -6,12 +6,11 @@ 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;
-    assert_eq!(immediate_dominators[0], None);
-    assert_eq!(immediate_dominators[1], Some(0));
-    assert_eq!(immediate_dominators[2], Some(0));
-    assert_eq!(immediate_dominators[3], Some(0));
+    let d = dominators(&graph);
+    assert_eq!(d.immediate_dominator(0), None);
+    assert_eq!(d.immediate_dominator(1), Some(0));
+    assert_eq!(d.immediate_dominator(2), Some(0));
+    assert_eq!(d.immediate_dominator(3), Some(0));
 }
 
 #[test]
@@ -22,15 +21,14 @@ 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;
-    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));
-    assert_eq!(immediate_dominators[3], Some(6));
-    assert_eq!(immediate_dominators[4], Some(6));
-    assert_eq!(immediate_dominators[5], Some(6));
-    assert_eq!(immediate_dominators[6], None);
+    let d = dominators(&graph);
+    assert_eq!(d.immediate_dominator(0), None); // <-- note that 0 is not in graph
+    assert_eq!(d.immediate_dominator(1), Some(6));
+    assert_eq!(d.immediate_dominator(2), Some(6));
+    assert_eq!(d.immediate_dominator(3), Some(6));
+    assert_eq!(d.immediate_dominator(4), Some(6));
+    assert_eq!(d.immediate_dominator(5), Some(6));
+    assert_eq!(d.immediate_dominator(6), None);
 }
 
 #[test]
@@ -47,11 +45,11 @@ fn paper_slt() {
 #[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 d = dominators(&graph);
+    assert_eq!(d.immediate_dominator(0), None);
+    assert_eq!(d.immediate_dominator(1), None);
+    assert_eq!(d.immediate_dominator(2), Some(1));
+    assert_eq!(d.immediate_dominator(3), Some(2));
 }
 
 #[test]
@@ -75,8 +73,7 @@ fn transitive_dominator() {
         ],
     );
 
-    let dom_tree = dominators(&graph);
-    let immediate_dominators = &dom_tree.immediate_dominators;
-    assert_eq!(immediate_dominators[2], Some(0));
-    assert_eq!(immediate_dominators[3], Some(0)); // This used to return Some(1).
+    let d = dominators(&graph);
+    assert_eq!(d.immediate_dominator(2), Some(0));
+    assert_eq!(d.immediate_dominator(3), Some(0)); // This used to return Some(1).
 }
diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index 34b6250a89d..f979f736b15 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -1562,6 +1562,23 @@ impl Location {
     }
 }
 
+/// `DefLocation` represents the location of a definition - either an argument or an assignment
+/// within MIR body.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum DefLocation {
+    Argument,
+    Body(Location),
+}
+
+impl DefLocation {
+    pub fn dominates(self, location: Location, dominators: &Dominators<BasicBlock>) -> bool {
+        match self {
+            DefLocation::Argument => true,
+            DefLocation::Body(def) => def.successor_within_block().dominates(location, dominators),
+        }
+    }
+}
+
 // Some nodes are used a lot. Make sure they don't unintentionally get bigger.
 #[cfg(all(target_arch = "x86_64", target_pointer_width = "64"))]
 mod size_asserts {
diff --git a/compiler/rustc_mir_transform/src/ssa.rs b/compiler/rustc_mir_transform/src/ssa.rs
index af9514ed6bb..fad58930e3a 100644
--- a/compiler/rustc_mir_transform/src/ssa.rs
+++ b/compiler/rustc_mir_transform/src/ssa.rs
@@ -15,7 +15,7 @@ use rustc_middle::mir::*;
 
 pub struct SsaLocals {
     /// Assignments to each local. This defines whether the local is SSA.
-    assignments: IndexVec<Local, Set1<LocationExtended>>,
+    assignments: IndexVec<Local, Set1<DefLocation>>,
     /// We visit the body in reverse postorder, to ensure each local is assigned before it is used.
     /// We remember the order in which we saw the assignments to compute the SSA values in a single
     /// pass.
@@ -27,55 +27,18 @@ pub struct SsaLocals {
     direct_uses: IndexVec<Local, u32>,
 }
 
-/// We often encounter MIR bodies with 1 or 2 basic blocks. In those cases, it's unnecessary to
-/// actually compute dominators, we can just compare block indices because bb0 is always the first
-/// block, and in any body all other blocks are always dominated by bb0.
-struct SmallDominators<'a> {
-    inner: Option<&'a Dominators<BasicBlock>>,
-}
-
-impl SmallDominators<'_> {
-    fn dominates(&self, first: Location, second: Location) -> bool {
-        if first.block == second.block {
-            first.statement_index <= second.statement_index
-        } else if let Some(inner) = &self.inner {
-            inner.dominates(first.block, second.block)
-        } else {
-            first.block < second.block
-        }
-    }
-
-    fn check_dominates(&mut self, set: &mut Set1<LocationExtended>, loc: Location) {
-        let assign_dominates = match *set {
-            Set1::Empty | Set1::Many => false,
-            Set1::One(LocationExtended::Arg) => true,
-            Set1::One(LocationExtended::Plain(assign)) => {
-                self.dominates(assign.successor_within_block(), loc)
-            }
-        };
-        // We are visiting a use that is not dominated by an assignment.
-        // Either there is a cycle involved, or we are reading for uninitialized local.
-        // Bail out.
-        if !assign_dominates {
-            *set = Set1::Many;
-        }
-    }
-}
-
 impl SsaLocals {
     pub fn new<'tcx>(body: &Body<'tcx>) -> SsaLocals {
         let assignment_order = Vec::with_capacity(body.local_decls.len());
 
         let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls);
-        let dominators =
-            if body.basic_blocks.len() > 2 { Some(body.basic_blocks.dominators()) } else { None };
-        let dominators = SmallDominators { inner: dominators };
+        let dominators = body.basic_blocks.dominators();
 
         let direct_uses = IndexVec::from_elem(0, &body.local_decls);
         let mut visitor = SsaVisitor { assignments, assignment_order, dominators, direct_uses };
 
         for local in body.args_iter() {
-            visitor.assignments[local] = Set1::One(LocationExtended::Arg);
+            visitor.assignments[local] = Set1::One(DefLocation::Argument);
         }
 
         // For SSA assignments, a RPO visit will see the assignment before it sees any use.
@@ -131,14 +94,7 @@ impl SsaLocals {
         location: Location,
     ) -> bool {
         match self.assignments[local] {
-            Set1::One(LocationExtended::Arg) => true,
-            Set1::One(LocationExtended::Plain(ass)) => {
-                if ass.block == location.block {
-                    ass.statement_index < location.statement_index
-                } else {
-                    dominators.dominates(ass.block, location.block)
-                }
-            }
+            Set1::One(def) => def.dominates(location, dominators),
             _ => false,
         }
     }
@@ -148,7 +104,7 @@ impl SsaLocals {
         body: &'a Body<'tcx>,
     ) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>, Location)> + 'a {
         self.assignment_order.iter().filter_map(|&local| {
-            if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local] {
+            if let Set1::One(DefLocation::Body(loc)) = self.assignments[local] {
                 // `loc` must point to a direct assignment to `local`.
                 let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
                 let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() };
@@ -166,7 +122,7 @@ impl SsaLocals {
         mut f: impl FnMut(Local, &mut Rvalue<'tcx>, Location),
     ) {
         for &local in &self.assignment_order {
-            if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local] {
+            if let Set1::One(DefLocation::Body(loc)) = self.assignments[local] {
                 // `loc` must point to a direct assignment to `local`.
                 let bbs = basic_blocks.as_mut_preserves_cfg();
                 let bb = &mut bbs[loc.block];
@@ -224,19 +180,29 @@ impl SsaLocals {
     }
 }
 
-#[derive(Copy, Clone, Debug, PartialEq, Eq)]
-enum LocationExtended {
-    Plain(Location),
-    Arg,
-}
-
 struct SsaVisitor<'a> {
-    dominators: SmallDominators<'a>,
-    assignments: IndexVec<Local, Set1<LocationExtended>>,
+    dominators: &'a Dominators<BasicBlock>,
+    assignments: IndexVec<Local, Set1<DefLocation>>,
     assignment_order: Vec<Local>,
     direct_uses: IndexVec<Local, u32>,
 }
 
+impl SsaVisitor<'_> {
+    fn check_dominates(&mut self, local: Local, loc: Location) {
+        let set = &mut self.assignments[local];
+        let assign_dominates = match *set {
+            Set1::Empty | Set1::Many => false,
+            Set1::One(def) => def.dominates(loc, self.dominators),
+        };
+        // We are visiting a use that is not dominated by an assignment.
+        // Either there is a cycle involved, or we are reading for uninitialized local.
+        // Bail out.
+        if !assign_dominates {
+            *set = Set1::Many;
+        }
+    }
+}
+
 impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
     fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) {
         match ctxt {
@@ -254,7 +220,7 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
                 self.assignments[local] = Set1::Many;
             }
             PlaceContext::NonMutatingUse(_) => {
-                self.dominators.check_dominates(&mut self.assignments[local], loc);
+                self.check_dominates(local, loc);
                 self.direct_uses[local] += 1;
             }
             PlaceContext::NonUse(_) => {}
@@ -269,7 +235,7 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
                 let new_ctxt = PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy);
 
                 self.visit_projection(place.as_ref(), new_ctxt, loc);
-                self.dominators.check_dominates(&mut self.assignments[place.local], loc);
+                self.check_dominates(place.local, loc);
             }
             return;
         } else {
@@ -280,7 +246,7 @@ impl<'tcx> Visitor<'tcx> for SsaVisitor<'_> {
 
     fn visit_assign(&mut self, place: &Place<'tcx>, rvalue: &Rvalue<'tcx>, loc: Location) {
         if let Some(local) = place.as_local() {
-            self.assignments[local].insert(LocationExtended::Plain(loc));
+            self.assignments[local].insert(DefLocation::Body(loc));
             if let Set1::One(_) = self.assignments[local] {
                 // Only record if SSA-like, to avoid growing the vector needlessly.
                 self.assignment_order.push(local);
@@ -356,7 +322,7 @@ fn compute_copy_classes(ssa: &mut SsaLocals, body: &Body<'_>) {
 #[derive(Debug)]
 pub(crate) struct StorageLiveLocals {
     /// Set of "StorageLive" statements for each local.
-    storage_live: IndexVec<Local, Set1<LocationExtended>>,
+    storage_live: IndexVec<Local, Set1<DefLocation>>,
 }
 
 impl StorageLiveLocals {
@@ -366,13 +332,13 @@ impl StorageLiveLocals {
     ) -> StorageLiveLocals {
         let mut storage_live = IndexVec::from_elem(Set1::Empty, &body.local_decls);
         for local in always_storage_live_locals.iter() {
-            storage_live[local] = Set1::One(LocationExtended::Arg);
+            storage_live[local] = Set1::One(DefLocation::Argument);
         }
         for (block, bbdata) in body.basic_blocks.iter_enumerated() {
             for (statement_index, statement) in bbdata.statements.iter().enumerate() {
                 if let StatementKind::StorageLive(local) = statement.kind {
                     storage_live[local]
-                        .insert(LocationExtended::Plain(Location { block, statement_index }));
+                        .insert(DefLocation::Body(Location { block, statement_index }));
                 }
             }
         }