From f16d2b1629d62db306bd9d03676524dc848c94b4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Mi=C4=85sko?= <tomasz.miasko@gmail.com>
Date: Sun, 14 May 2023 00:00:00 +0000
Subject: [PATCH] Start node has no immediate dominator

Change the immediate_dominator return type to Option, and use None to
indicate that node has no immediate dominator.

Also fix the issue where the start node would be returned as its own
immediate dominator.
---
 .../src/transform/validate.rs                 |  2 +-
 .../src/graph/dominators/mod.rs               | 22 +++++++++----------
 .../src/graph/dominators/tests.rs             | 14 ++++++++++--
 3 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
index c50e937d84f..3c350e25ba6 100644
--- a/compiler/rustc_const_eval/src/transform/validate.rs
+++ b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -164,7 +164,7 @@ fn check_cleanup_control_flow(&self) {
                 if let Some(root) = post_contract_node.get(&bb) {
                     break *root;
                 }
-                let parent = doms.immediate_dominator(bb);
+                let parent = doms.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 e76bdac2864..a7de709ba72 100644
--- a/compiler/rustc_data_structures/src/graph/dominators/mod.rs
+++ b/compiler/rustc_data_structures/src/graph/dominators/mod.rs
@@ -242,7 +242,9 @@ pub fn dominators<G: ControlFlowGraph>(graph: G) -> Dominators<G::Node> {
         immediate_dominators[*node] = Some(pre_order_to_real[idom[idx]]);
     }
 
-    Dominators { post_order_rank, immediate_dominators }
+    let start_node = graph.start_node();
+    immediate_dominators[start_node] = None;
+    Dominators { start_node, post_order_rank, immediate_dominators }
 }
 
 /// Evaluate the link-eval virtual forest, providing the currently minimum semi
@@ -308,6 +310,7 @@ fn compress(
 /// Tracks the list of dominators for each node.
 #[derive(Clone, Debug)]
 pub struct Dominators<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
@@ -316,14 +319,14 @@ pub struct Dominators<N: Idx> {
 }
 
 impl<Node: Idx> Dominators<Node> {
-    /// Whether the given Node has an immediate dominator.
+    /// Returns true if node is reachable from the start node.
     pub fn is_reachable(&self, node: Node) -> bool {
-        self.immediate_dominators[node].is_some()
+        node == self.start_node || self.immediate_dominators[node].is_some()
     }
 
-    pub fn immediate_dominator(&self, node: Node) -> Node {
-        assert!(self.is_reachable(node), "node {node:?} is not reachable");
-        self.immediate_dominators[node].unwrap()
+    /// Returns the immediate dominator of node, if any.
+    pub fn immediate_dominator(&self, node: Node) -> Option<Node> {
+        self.immediate_dominators[node]
     }
 
     /// Provides an iterator over each dominator up the CFG, for the given Node.
@@ -357,12 +360,7 @@ impl<'dom, Node: Idx> Iterator for Iter<'dom, Node> {
 
     fn next(&mut self) -> Option<Self::Item> {
         if let Some(node) = self.node {
-            let dom = self.dominators.immediate_dominator(node);
-            if dom == node {
-                self.node = None; // reached the root
-            } else {
-                self.node = Some(dom);
-            }
+            self.node = self.dominators.immediate_dominator(node);
             Some(node)
         } else {
             None
diff --git a/compiler/rustc_data_structures/src/graph/dominators/tests.rs b/compiler/rustc_data_structures/src/graph/dominators/tests.rs
index ff31d8f7fdc..8b124516623 100644
--- a/compiler/rustc_data_structures/src/graph/dominators/tests.rs
+++ b/compiler/rustc_data_structures/src/graph/dominators/tests.rs
@@ -8,7 +8,7 @@ fn diamond() {
 
     let dominators = dominators(&graph);
     let immediate_dominators = &dominators.immediate_dominators;
-    assert_eq!(immediate_dominators[0], Some(0));
+    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));
@@ -30,7 +30,7 @@ fn paper() {
     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], Some(6));
+    assert_eq!(immediate_dominators[6], None);
 }
 
 #[test]
@@ -43,3 +43,13 @@ fn paper_slt() {
 
     dominators(&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));
+}