Rollup merge of #73133 - doctorn:unwind-mir-validation, r=jonas-schievink

Enforce unwind invariants

I had a quick look at #72959. The failure message probably needs to be more detailed but I just wanted to check I got the right idea. I have no idea how to right a test for this either...

r? @jonas-schievink

Resolves #72959 (hypothetically)

Author: Dylan-DPC

src/librustc_mir/transform/validate.rs

@@ -11,6 +11,12 @@ use rustc_middle::{
 };
 use rustc_span::def_id::DefId;
 
+#[derive(Copy, Clone, Debug)]
+enum EdgeKind {
+    Unwind,
+    Normal,
+}
+
 pub struct Validator {
     /// Describes at which point in the pipeline this validation is happening.
     pub when: String,
@@ -49,8 +55,31 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
         );
     }
 
-    fn check_bb(&self, location: Location, bb: BasicBlock) {
-        if self.body.basic_blocks().get(bb).is_none() {
+    fn check_edge(&self, location: Location, bb: BasicBlock, edge_kind: EdgeKind) {
+        if let Some(bb) = self.body.basic_blocks().get(bb) {
+            let src = self.body.basic_blocks().get(location.block).unwrap();
+            match (src.is_cleanup, bb.is_cleanup, edge_kind) {
+                // Non-cleanup blocks can jump to non-cleanup blocks along non-unwind edges
+                (false, false, EdgeKind::Normal)
+                // Non-cleanup blocks can jump to cleanup blocks along unwind edges
+                | (false, true, EdgeKind::Unwind)
+                // Cleanup blocks can jump to cleanup blocks along non-unwind edges
+                | (true, true, EdgeKind::Normal) => {}
+                // All other jumps are invalid
+                _ => {
+                    self.fail(
+                        location,
+                        format!(
+                            "{:?} edge to {:?} violates unwind invariants (cleanup {:?} -> {:?})",
+                            edge_kind,
+                            bb,
+                            src.is_cleanup,
+                            bb.is_cleanup,
+                        )
+                    )
+                }
+            }
+        } else {
             self.fail(location, format!("encountered jump to invalid basic block {:?}", bb))
         }
     }
@@ -92,7 +121,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
     fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
         match &terminator.kind {
             TerminatorKind::Goto { target } => {
-                self.check_bb(location, *target);
+                self.check_edge(location, *target, EdgeKind::Normal);
             }
             TerminatorKind::SwitchInt { targets, values, .. } => {
                 if targets.len() != values.len() + 1 {
@@ -106,19 +135,19 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     );
                 }
                 for target in targets {
-                    self.check_bb(location, *target);
+                    self.check_edge(location, *target, EdgeKind::Normal);
                 }
             }
             TerminatorKind::Drop { target, unwind, .. } => {
-                self.check_bb(location, *target);
+                self.check_edge(location, *target, EdgeKind::Normal);
                 if let Some(unwind) = unwind {
-                    self.check_bb(location, *unwind);
+                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                 }
             }
             TerminatorKind::DropAndReplace { target, unwind, .. } => {
-                self.check_bb(location, *target);
+                self.check_edge(location, *target, EdgeKind::Normal);
                 if let Some(unwind) = unwind {
-                    self.check_bb(location, *unwind);
+                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                 }
             }
             TerminatorKind::Call { func, destination, cleanup, .. } => {
@@ -131,10 +160,10 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     ),
                 }
                 if let Some((_, target)) = destination {
-                    self.check_bb(location, *target);
+                    self.check_edge(location, *target, EdgeKind::Normal);
                 }
                 if let Some(cleanup) = cleanup {
-                    self.check_bb(location, *cleanup);
+                    self.check_edge(location, *cleanup, EdgeKind::Unwind);
                 }
             }
             TerminatorKind::Assert { cond, target, cleanup, .. } => {
@@ -148,30 +177,30 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                         ),
                     );
                 }
-                self.check_bb(location, *target);
+                self.check_edge(location, *target, EdgeKind::Normal);
                 if let Some(cleanup) = cleanup {
-                    self.check_bb(location, *cleanup);
+                    self.check_edge(location, *cleanup, EdgeKind::Unwind);
                 }
             }
             TerminatorKind::Yield { resume, drop, .. } => {
-                self.check_bb(location, *resume);
+                self.check_edge(location, *resume, EdgeKind::Normal);
                 if let Some(drop) = drop {
-                    self.check_bb(location, *drop);
+                    self.check_edge(location, *drop, EdgeKind::Normal);
                 }
             }
             TerminatorKind::FalseEdge { real_target, imaginary_target } => {
-                self.check_bb(location, *real_target);
-                self.check_bb(location, *imaginary_target);
+                self.check_edge(location, *real_target, EdgeKind::Normal);
+                self.check_edge(location, *imaginary_target, EdgeKind::Normal);
             }
             TerminatorKind::FalseUnwind { real_target, unwind } => {
-                self.check_bb(location, *real_target);
+                self.check_edge(location, *real_target, EdgeKind::Normal);
                 if let Some(unwind) = unwind {
-                    self.check_bb(location, *unwind);
+                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                 }
             }
             TerminatorKind::InlineAsm { destination, .. } => {
                 if let Some(destination) = destination {
-                    self.check_bb(location, *destination);
+                    self.check_edge(location, *destination, EdgeKind::Normal);
                 }
             }
             // Nothing to validate for these.