Add cycle checking to cleanup control flow validation

2 files changed, 54 insertions, 22 deletions

diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
index 9f429d3a7d9..b4f1ab62267 100644
--- a/compiler/rustc_const_eval/src/transform/validate.rs
+++ b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -1,9 +1,8 @@
 //! Validates the MIR to ensure that invariants are upheld.
 
-use std::collections::hash_map::Entry;
-
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_index::bit_set::BitSet;
+use rustc_index::vec::IndexVec;
 use rustc_infer::traits::Reveal;
 use rustc_middle::mir::interpret::Scalar;
 use rustc_middle::mir::visit::NonUseContext::VarDebugInfo;
@@ -140,23 +139,27 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
     fn check_cleanup_control_flow(&self) {
         let doms = self.body.basic_blocks.dominators();
         let mut post_contract_node = FxHashMap::default();
+        // Reusing the allocation across invocations of the closure
+        let mut dom_path = vec![];
         let mut get_post_contract_node = |mut bb| {
-            if let Some(res) = post_contract_node.get(&bb) {
-                return *res;
-            }
-            let mut dom_path = vec![];
-            while self.body.basic_blocks[bb].is_cleanup {
+            let root = loop {
+                if let Some(root) = post_contract_node.get(&bb) {
+                    break *root;
+                }
+                let parent = doms.immediate_dominator(bb);
                 dom_path.push(bb);
-                bb = doms.immediate_dominator(bb);
-            }
-            let root = *dom_path.last().unwrap();
-            for bb in dom_path {
+                if !self.body.basic_blocks[parent].is_cleanup {
+                    break bb;
+                }
+                bb = parent;
+            };
+            for bb in dom_path.drain(..) {
                 post_contract_node.insert(bb, root);
             }
             root
         };
 
-        let mut parent = FxHashMap::default();
+        let mut parent = IndexVec::from_elem(None, &self.body.basic_blocks);
         for (bb, bb_data) in self.body.basic_blocks.iter_enumerated() {
             if !bb_data.is_cleanup || !self.reachable_blocks.contains(bb) {
                 continue;
@@ -167,23 +170,49 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 if s == bb {
                     continue;
                 }
-                match parent.entry(bb) {
-                    Entry::Vacant(e) => {
-                        e.insert(s);
+                let parent = &mut parent[bb];
+                match parent {
+                    None => {
+                        *parent = Some(s);
                     }
-                    Entry::Occupied(e) if s != *e.get() => self.fail(
+                    Some(e) if *e == s => (),
+                    Some(e) => self.fail(
                         Location { block: bb, statement_index: 0 },
                         format!(
                             "Cleanup control flow violation: The blocks dominated by {:?} have edges to both {:?} and {:?}",
                             bb,
                             s,
-                            *e.get()
+                            *e
                         )
                     ),
-                    Entry::Occupied(_) => (),
                 }
             }
         }
+
+        // Check for cycles
+        let mut stack = FxHashSet::default();
+        for i in 0..parent.len() {
+            let mut bb = BasicBlock::from_usize(i);
+            stack.clear();
+            stack.insert(bb);
+            loop {
+                let Some(parent )= parent[bb].take() else {
+                    break
+                };
+                let no_cycle = stack.insert(parent);
+                if !no_cycle {
+                    self.fail(
+                        Location { block: bb, statement_index: 0 },
+                        format!(
+                            "Cleanup control flow violation: Cycle involving edge {:?} -> {:?}",
+                            bb, parent,
+                        ),
+                    );
+                    break;
+                }
+                bb = parent;
+            }
+        }
     }
 
     /// Check if src can be assigned into dest.
diff --git a/compiler/rustc_middle/src/mir/syntax.rs b/compiler/rustc_middle/src/mir/syntax.rs
index 0c395cae566..52c2b10cbbe 100644
--- a/compiler/rustc_middle/src/mir/syntax.rs
+++ b/compiler/rustc_middle/src/mir/syntax.rs
@@ -512,13 +512,16 @@ pub struct CopyNonOverlapping<'tcx> {
 ///     must also be `cleanup`. This is a part of the type system and checked statically, so it is
 ///     still an error to have such an edge in the CFG even if it's known that it won't be taken at
 ///     runtime.
-///  4. The induced subgraph on cleanup blocks must look roughly like an upside down tree. This is
-///     necessary to ensure that landing pad information can be correctly codegened. More precisely:
+///  4. The control flow between cleanup blocks must look like an upside down tree. Roughly
+///     speaking, this means that control flow that looks like a V is allowed, while control flow
+///     that looks like a W is not. This is necessary to ensure that landing pad information can be
+///     correctly codegened on MSVC. More precisely:
 ///
 ///     Begin with the standard control flow graph `G`. Modify `G` as follows: for any two cleanup
 ///     vertices `u` and `v` such that `u` dominates `v`, contract `u` and `v` into a single vertex,
-///     deleting self edges and duplicate edges in the process. The cleanup blocks of the resulting
-///     graph must form an inverted forest.
+///     deleting self edges and duplicate edges in the process. Now remove all vertices from `G`
+///     that are not cleanup vertices or are not reachable. The resulting graph must be an inverted
+///     tree, that is each vertex may have at most one successor and there may be no cycles.
 #[derive(Clone, TyEncodable, TyDecodable, Hash, HashStable, PartialEq, TypeFoldable, TypeVisitable)]
 pub enum TerminatorKind<'tcx> {
     /// Block has one successor; we continue execution there.