Auto merge of #66329 - ktrianta:mir-opt-unreachable-propagation, r=oli-obk

Add unreachable propagation mir optimization pass

@oli-obk suggested we create a MIR pass that optimizes away basic blocks that lead only to basic blocks with terminator kind **unreachable**. This is a first take on this, which we started with @gilescope at RustFest Impl Days.

The test currently fails when the compiled program runs (undefined behaviour). Is there a way to avoid running the compiled program?

8 files changed, 449 insertions, 67 deletions

diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
index 22aed9a9e66..3c37eccc184 100644
--- a/src/librustc_mir/transform/mod.rs
+++ b/src/librustc_mir/transform/mod.rs
@@ -36,6 +36,7 @@ pub mod simplify;
 pub mod simplify_branches;
 pub mod simplify_try;
 pub mod uninhabited_enum_branching;
+pub mod unreachable_prop;
 
 pub(crate) fn provide(providers: &mut Providers<'_>) {
     self::check_unsafety::provide(providers);
@@ -299,6 +300,7 @@ fn run_optimization_passes<'tcx>(
             // From here on out, regions are gone.
             &erase_regions::EraseRegions,
             // Optimizations begin.
+            &unreachable_prop::UnreachablePropagation,
             &uninhabited_enum_branching::UninhabitedEnumBranching,
             &simplify::SimplifyCfg::new("after-uninhabited-enum-branching"),
             &inline::Inline,
diff --git a/src/librustc_mir/transform/unreachable_prop.rs b/src/librustc_mir/transform/unreachable_prop.rs
new file mode 100644
index 00000000000..27173e0c171
--- /dev/null
+++ b/src/librustc_mir/transform/unreachable_prop.rs
@@ -0,0 +1,108 @@
+//! A pass that propagates the unreachable terminator of a block to its predecessors
+//! when all of their successors are unreachable. This is achieved through a
+//! post-order traversal of the blocks.
+
+use crate::transform::simplify;
+use crate::transform::{MirPass, MirSource};
+use rustc::mir::*;
+use rustc::ty::TyCtxt;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use std::borrow::Cow;
+
+pub struct UnreachablePropagation;
+
+impl MirPass<'_> for UnreachablePropagation {
+    fn run_pass<'tcx>(&self, tcx: TyCtxt<'tcx>, _: MirSource<'tcx>, body: &mut BodyAndCache<'tcx>) {
+        if tcx.sess.opts.debugging_opts.mir_opt_level < 3 {
+            // Enable only under -Zmir-opt-level=3 as in some cases (check the deeply-nested-opt
+            // perf benchmark) LLVM may spend quite a lot of time optimizing the generated code.
+            return;
+        }
+
+        let mut unreachable_blocks = FxHashSet::default();
+        let mut replacements = FxHashMap::default();
+
+        for (bb, bb_data) in traversal::postorder(body) {
+            let terminator = bb_data.terminator();
+            // HACK: If the block contains any asm statement it is not regarded as unreachable.
+            // This is a temporary solution that handles possibly diverging asm statements.
+            // Accompanying testcases: mir-opt/unreachable_asm.rs and mir-opt/unreachable_asm_2.rs
+            let asm_stmt_in_block = || {
+                bb_data.statements.iter().any(|stmt: &Statement<'_>| match stmt.kind {
+                    StatementKind::InlineAsm(..) => true,
+                    _ => false,
+                })
+            };
+
+            if terminator.kind == TerminatorKind::Unreachable && !asm_stmt_in_block() {
+                unreachable_blocks.insert(bb);
+            } else {
+                let is_unreachable = |succ: BasicBlock| unreachable_blocks.contains(&succ);
+                let terminator_kind_opt = remove_successors(&terminator.kind, is_unreachable);
+
+                if let Some(terminator_kind) = terminator_kind_opt {
+                    if terminator_kind == TerminatorKind::Unreachable && !asm_stmt_in_block() {
+                        unreachable_blocks.insert(bb);
+                    }
+                    replacements.insert(bb, terminator_kind);
+                }
+            }
+        }
+
+        let replaced = !replacements.is_empty();
+        for (bb, terminator_kind) in replacements {
+            body.basic_blocks_mut()[bb].terminator_mut().kind = terminator_kind;
+        }
+
+        if replaced {
+            simplify::remove_dead_blocks(body);
+        }
+    }
+}
+
+fn remove_successors<F>(
+    terminator_kind: &TerminatorKind<'tcx>,
+    predicate: F,
+) -> Option<TerminatorKind<'tcx>>
+where
+    F: Fn(BasicBlock) -> bool,
+{
+    match *terminator_kind {
+        TerminatorKind::Goto { target } if predicate(target) => Some(TerminatorKind::Unreachable),
+        TerminatorKind::SwitchInt { ref discr, switch_ty, ref values, ref targets } => {
+            let original_targets_len = targets.len();
+            let (otherwise, targets) = targets.split_last().unwrap();
+            let retained = values
+                .iter()
+                .zip(targets.iter())
+                .filter(|(_, &t)| !predicate(t))
+                .collect::<Vec<_>>();
+            let mut values = retained.iter().map(|&(v, _)| *v).collect::<Vec<_>>();
+            let mut targets = retained.iter().map(|&(_, d)| *d).collect::<Vec<_>>();
+
+            if !predicate(*otherwise) {
+                targets.push(*otherwise);
+            } else {
+                values.pop();
+            }
+
+            let retained_targets_len = targets.len();
+
+            if targets.is_empty() {
+                Some(TerminatorKind::Unreachable)
+            } else if targets.len() == 1 {
+                Some(TerminatorKind::Goto { target: targets[0] })
+            } else if original_targets_len != retained_targets_len {
+                Some(TerminatorKind::SwitchInt {
+                    discr: discr.clone(),
+                    switch_ty,
+                    values: Cow::from(values),
+                    targets,
+                })
+            } else {
+                None
+            }
+        }
+        _ => None,
+    }
+}
diff --git a/src/test/mir-opt/simplify_try.rs b/src/test/mir-opt/simplify_try.rs
index 656b405ef34..d85eff45b49 100644
--- a/src/test/mir-opt/simplify_try.rs
+++ b/src/test/mir-opt/simplify_try.rs
@@ -47,25 +47,22 @@ fn main() {
 //     }
 //     bb0: {
 //         _5 = discriminant(_1);
-//         switchInt(move _5) -> [0isize: bb4, 1isize: bb2, otherwise: bb1];
+//         switchInt(move _5) -> [0isize: bb3, otherwise: bb1];
 //     }
 //     bb1: {
-//         unreachable;
-//     }
-//     bb2: {
 //         _6 = ((_1 as Err).0: i32);
 //         ((_0 as Err).0: i32) = move _6;
 //         discriminant(_0) = 1;
-//         goto -> bb3;
+//         goto -> bb2;
 //     }
-//     bb3: {
+//     bb2: {
 //         return;
 //     }
-//     bb4: {
+//     bb3: {
 //         _10 = ((_1 as Ok).0: u32);
 //         ((_0 as Ok).0: u32) = move _10;
 //         discriminant(_0) = 0;
-//         goto -> bb3;
+//         goto -> bb2;
 //     }
 // }
 // END rustc.try_identity.SimplifyArmIdentity.before.mir
@@ -109,25 +106,22 @@ fn main() {
 //     }
 //     bb0: {
 //         _5 = discriminant(_1);
-//         switchInt(move _5) -> [0isize: bb4, 1isize: bb2, otherwise: bb1];
+//         switchInt(move _5) -> [0isize: bb3, otherwise: bb1];
 //     }
 //     bb1: {
-//         unreachable;
-//     }
-//     bb2: {
 //         _0 = move _1;
 //         nop;
 //         nop;
-//         goto -> bb3;
+//         goto -> bb2;
 //     }
-//     bb3: {
+//     bb2: {
 //         return;
 //     }
-//     bb4: {
+//     bb3: {
 //         _0 = move _1;
 //         nop;
 //         nop;
-//         goto -> bb3;
+//         goto -> bb2;
 //     }
 // }
 // END rustc.try_identity.SimplifyArmIdentity.after.mir
diff --git a/src/test/mir-opt/uninhabited_enum_branching.rs b/src/test/mir-opt/uninhabited_enum_branching.rs
index aa56918a9b8..dda5fd4fb75 100644
--- a/src/test/mir-opt/uninhabited_enum_branching.rs
+++ b/src/test/mir-opt/uninhabited_enum_branching.rs
@@ -45,53 +45,47 @@ fn main() {
 //   StorageLive(_2);
 //   _2 = Test1::C;
 //   _3 = discriminant(_2);
-//   switchInt(move _3) -> [0isize: bb3, 1isize: bb4, 2isize: bb1, otherwise: bb2];
+//   switchInt(move _3) -> [0isize: bb2, 1isize: bb3, otherwise: bb1];
 // }
 // bb1: {
 //   StorageLive(_5);
 //   _5 = const "C";
 //   _1 = &(*_5);
 //   StorageDead(_5);
-//   goto -> bb5;
+//   goto -> bb4;
 // }
 // bb2: {
-//   unreachable;
-// }
-// bb3: {
 //   _1 = const "A(Empty)";
-//   goto -> bb5;
+//   goto -> bb4;
 // }
-// bb4: {
+// bb3: {
 //   StorageLive(_4);
 //   _4 = const "B(Empty)";
 //   _1 = &(*_4);
 //   StorageDead(_4);
-//   goto -> bb5;
+//   goto -> bb4;
 // }
-// bb5: {
+// bb4: {
 //   StorageDead(_2);
 //   StorageDead(_1);
 //   StorageLive(_6);
 //   StorageLive(_7);
 //   _7 = Test2::D;
 //   _8 = discriminant(_7);
-//   switchInt(move _8) -> [4isize: bb8, 5isize: bb6, otherwise: bb7];
+//   switchInt(move _8) -> [4isize: bb6, otherwise: bb5];
 // }
-// bb6: {
+// bb5: {
 //   StorageLive(_9);
 //   _9 = const "E";
 //   _6 = &(*_9);
 //   StorageDead(_9);
-//   goto -> bb9;
-// }
-// bb7: {
-//   unreachable;
+//   goto -> bb7;
 // }
-// bb8: {
+// bb6: {
 //   _6 = const "D";
-//   goto -> bb9;
+//   goto -> bb7;
 // }
-// bb9: {
+// bb7: {
 //   StorageDead(_7);
 //   StorageDead(_6);
 //   _0 = ();
@@ -114,53 +108,47 @@ fn main() {
 //   StorageLive(_2);
 //   _2 = Test1::C;
 //   _3 = discriminant(_2);
-//   switchInt(move _3) -> [2isize: bb1, otherwise: bb2];
+//   switchInt(move _3) -> bb1;
 // }
 // bb1: {
 //   StorageLive(_5);
 //   _5 = const "C";
 //   _1 = &(*_5);
 //   StorageDead(_5);
-//   goto -> bb5;
+//   goto -> bb4;
 // }
 // bb2: {
-//   unreachable;
-// }
-// bb3: {
 //   _1 = const "A(Empty)";
-//   goto -> bb5;
+//   goto -> bb4;
 // }
-// bb4: {
+// bb3: {
 //   StorageLive(_4);
 //   _4 = const "B(Empty)";
 //   _1 = &(*_4);
 //   StorageDead(_4);
-//   goto -> bb5;
+//   goto -> bb4;
 // }
-// bb5: {
+// bb4: {
 //   StorageDead(_2);
 //   StorageDead(_1);
 //   StorageLive(_6);
 //   StorageLive(_7);
 //   _7 = Test2::D;
 //   _8 = discriminant(_7);
-//   switchInt(move _8) -> [4isize: bb8, 5isize: bb6, otherwise: bb7];
+//   switchInt(move _8) -> [4isize: bb6, otherwise: bb5];
 // }
-// bb6: {
+// bb5: {
 //   StorageLive(_9);
 //   _9 = const "E";
 //   _6 = &(*_9);
 //   StorageDead(_9);
-//   goto -> bb9;
+//   goto -> bb7;
 // }
-// bb7: {
-//   unreachable;
-// }
-// bb8: {
+// bb6: {
 //   _6 = const "D";
-//   goto -> bb9;
+//   goto -> bb7;
 // }
-// bb9: {
+// bb7: {
 //   StorageDead(_7);
 //   StorageDead(_6);
 //   _0 = ();
@@ -183,9 +171,6 @@ fn main() {
 //   StorageLive(_2);
 //   _2 = Test1::C;
 //   _3 = discriminant(_2);
-//   switchInt(move _3) -> [2isize: bb1, otherwise: bb2];
-// }
-// bb1: {
 //   StorageLive(_5);
 //   _5 = const "C";
 //   _1 = &(*_5);
@@ -196,26 +181,20 @@ fn main() {
 //   StorageLive(_7);
 //   _7 = Test2::D;
 //   _8 = discriminant(_7);
-//   switchInt(move _8) -> [4isize: bb5, 5isize: bb3, otherwise: bb4];
-// }
-// bb2: {
-//   unreachable;
+//   switchInt(move _8) -> [4isize: bb2, otherwise: bb1];
 // }
-// bb3: {
+// bb1: {
 //   StorageLive(_9);
 //   _9 = const "E";
 //   _6 = &(*_9);
 //   StorageDead(_9);
-//   goto -> bb6;
-// }
-// bb4: {
-//   unreachable;
+//   goto -> bb3;
 // }
-// bb5: {
+// bb2: {
 //   _6 = const "D";
-//   goto -> bb6;
+//   goto -> bb3;
 // }
-// bb6: {
+// bb3: {
 //   StorageDead(_7);
 //   StorageDead(_6);
 //   _0 = ();
diff --git a/src/test/mir-opt/unreachable.rs b/src/test/mir-opt/unreachable.rs
new file mode 100644
index 00000000000..fa5c1a074ee
--- /dev/null
+++ b/src/test/mir-opt/unreachable.rs
@@ -0,0 +1,78 @@
+enum Empty {}
+
+fn empty() -> Option<Empty> {
+    None
+}
+
+fn main() {
+    if let Some(_x) = empty() {
+        let mut _y;
+
+        if true {
+            _y = 21;
+        } else {
+            _y = 42;
+        }
+
+        match _x { }
+    }
+}
+
+// END RUST SOURCE
+// START rustc.main.UnreachablePropagation.before.mir
+//      bb0: {
+//          StorageLive(_1);
+//          _1 = const empty() -> bb1;
+//      }
+//      bb1: {
+//          _2 = discriminant(_1);
+//          switchInt(move _2) -> [1isize: bb3, otherwise: bb2];
+//      }
+//      bb2: {
+//          _0 = ();
+//          StorageDead(_1);
+//          return;
+//      }
+//      bb3: {
+//          StorageLive(_3);
+//          _3 = move ((_1 as Some).0: Empty);
+//          StorageLive(_4);
+//          StorageLive(_5);
+//          StorageLive(_6);
+//          _6 = const true;
+//          switchInt(_6) -> [false: bb4, otherwise: bb5];
+//      }
+//      bb4: {
+//          _4 = const 42i32;
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//      bb5: {
+//          _4 = const 21i32;
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//      bb6: {
+//          StorageDead(_6);
+//          StorageDead(_5);
+//          StorageLive(_7);
+//          unreachable;
+//      }
+//  }
+// END rustc.main.UnreachablePropagation.before.mir
+// START rustc.main.UnreachablePropagation.after.mir
+//      bb0: {
+//          StorageLive(_1);
+//          _1 = const empty() -> bb1;
+//      }
+//      bb1: {
+//          _2 = discriminant(_1);
+//          goto -> bb2;
+//      }
+//      bb2: {
+//          _0 = ();
+//          StorageDead(_1);
+//          return;
+//      }
+//  }
+// END rustc.main.UnreachablePropagation.after.mir
diff --git a/src/test/mir-opt/unreachable_asm.rs b/src/test/mir-opt/unreachable_asm.rs
new file mode 100644
index 00000000000..ca614ac32b7
--- /dev/null
+++ b/src/test/mir-opt/unreachable_asm.rs
@@ -0,0 +1,72 @@
+// ignore-tidy-linelength
+#![feature(asm)]
+
+enum Empty {}
+
+fn empty() -> Option<Empty> {
+    None
+}
+
+fn main() {
+    if let Some(_x) = empty() {
+        let mut _y;
+
+        if true {
+            _y = 21;
+        } else {
+            _y = 42;
+        }
+
+        // asm instruction stops unreachable propagation to if else blocks bb4 and bb5.
+        unsafe { asm!("NOP"); }
+        match _x { }
+    }
+}
+
+// END RUST SOURCE
+// START rustc.main.UnreachablePropagation.before.mir
+//      bb4: {
+//          _4 = const 42i32;
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//      bb5: {
+//          _4 = const 21i32;
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//      bb6: {
+//          StorageDead(_6);
+//          StorageDead(_5);
+//          StorageLive(_7);
+//          asm!(InlineAsmInner { asm: "NOP", asm_str_style: Cooked, outputs: [], inputs: [], clobbers: [], volatile: true, alignstack: false, dialect: Att } : [] : []);
+//          _7 = ();
+//          StorageDead(_7);
+//          StorageLive(_8);
+//          unreachable;
+//      }
+//  }
+// END rustc.main.UnreachablePropagation.before.mir
+// START rustc.main.UnreachablePropagation.after.mir
+//      bb4: {
+//          _4 = const 42i32;
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//      bb5: {
+//          _4 = const 21i32;
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//      bb6: {
+//          StorageDead(_6);
+//          StorageDead(_5);
+//          StorageLive(_7);
+//          asm!(InlineAsmInner { asm: "NOP", asm_str_style: Cooked, outputs: [], inputs: [], clobbers: [], volatile: true, alignstack: false, dialect: Att } : [] : []);
+//          _7 = ();
+//          StorageDead(_7);
+//          StorageLive(_8);
+//          unreachable;
+//      }
+//  }
+// END rustc.main.UnreachablePropagation.after.mir
diff --git a/src/test/mir-opt/unreachable_asm_2.rs b/src/test/mir-opt/unreachable_asm_2.rs
new file mode 100644
index 00000000000..8fdbcfb5cab
--- /dev/null
+++ b/src/test/mir-opt/unreachable_asm_2.rs
@@ -0,0 +1,84 @@
+// ignore-tidy-linelength
+#![feature(asm)]
+
+enum Empty {}
+
+fn empty() -> Option<Empty> {
+    None
+}
+
+fn main() {
+    if let Some(_x) = empty() {
+        let mut _y;
+
+        if true {
+            // asm instruction stops unreachable propagation to block bb3.
+            unsafe { asm!("NOP"); }
+            _y = 21;
+        } else {
+            // asm instruction stops unreachable propagation to block bb3.
+            unsafe { asm!("NOP"); }
+            _y = 42;
+        }
+
+        match _x { }
+    }
+}
+
+// END RUST SOURCE
+// START rustc.main.UnreachablePropagation.before.mir
+//      bb3: {
+//          ...
+//          switchInt(_6) -> [false: bb4, otherwise: bb5];
+//      }
+//      bb4: {
+//          StorageLive(_8);
+//          asm!(InlineAsmInner { asm: "NOP", asm_str_style: Cooked, outputs: [], inputs: [], clobbers: [], volatile: true, alignstack: false, dialect: Att } : [] : []);
+//          _8 = ();
+//          StorageDead(_8);
+//          _4 = const 42i32;
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//          bb5: {
+//          StorageLive(_7);
+//          asm!(InlineAsmInner { asm: "NOP", asm_str_style: Cooked, outputs: [], inputs: [], clobbers: [], volatile: true, alignstack: false, dialect: Att } : [] : []);
+//          _7 = ();
+//          StorageDead(_7);
+//          _4 = const 21i32;
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//      bb6: {
+//          StorageDead(_6);
+//          StorageDead(_5);
+//          StorageLive(_9);
+//          unreachable;
+//      }
+//  }
+// END rustc.main.UnreachablePropagation.before.mir
+// START rustc.main.UnreachablePropagation.after.mir
+//      bb3: {
+//          ...
+//          switchInt(_6) -> [false: bb4, otherwise: bb5];
+//      }
+//      bb4: {
+//          StorageLive(_8);
+//          asm!(InlineAsmInner { asm: "NOP", asm_str_style: Cooked, outputs: [], inputs: [], clobbers: [], volatile: true, alignstack: false, dialect: Att } : [] : []);
+//          _8 = ();
+//          StorageDead(_8);
+//          _4 = const 42i32;
+//          _5 = ();
+//          unreachable;
+//      }
+//          bb5: {
+//          StorageLive(_7);
+//          asm!(InlineAsmInner { asm: "NOP", asm_str_style: Cooked, outputs: [], inputs: [], clobbers: [], volatile: true, alignstack: false, dialect: Att } : [] : []);
+//          _7 = ();
+//          StorageDead(_7);
+//          _4 = const 21i32;
+//          _5 = ();
+//          unreachable;
+//      }
+//  }
+// END rustc.main.UnreachablePropagation.after.mir
diff --git a/src/test/mir-opt/unreachable_diverging.rs b/src/test/mir-opt/unreachable_diverging.rs
new file mode 100644
index 00000000000..bf05019d5ce
--- /dev/null
+++ b/src/test/mir-opt/unreachable_diverging.rs
@@ -0,0 +1,65 @@
+pub enum Empty {}
+
+fn empty() -> Option<Empty> {
+    None
+}
+
+fn loop_forever() {
+    loop {}
+}
+
+fn main() {
+    let x = true;
+    if let Some(bomb) = empty() {
+        if x {
+            loop_forever()
+        }
+        match bomb {}
+    }
+}
+
+// END RUST SOURCE
+// START rustc.main.UnreachablePropagation.before.mir
+//      bb3: {
+//          StorageLive(_4);
+//          _4 = move ((_2 as Some).0: Empty);
+//          StorageLive(_5);
+//          StorageLive(_6);
+//          _6 = _1;
+//          switchInt(_6) -> [false: bb4, otherwise: bb5];
+//      }
+//      bb4: {
+//          _5 = ();
+//          goto -> bb6;
+//      }
+//      bb5: {
+//          _5 = const loop_forever() -> bb6;
+//      }
+//      bb6: {
+//          StorageDead(_6);
+//          StorageDead(_5);
+//          StorageLive(_7);
+//          unreachable;
+//      }
+//  }
+// END rustc.main.UnreachablePropagation.before.mir
+// START rustc.main.UnreachablePropagation.after.mir
+//      bb3: {
+//          StorageLive(_4);
+//          _4 = move ((_2 as Some).0: Empty);
+//          StorageLive(_5);
+//          StorageLive(_6);
+//          _6 = _1;
+//          goto -> bb4;
+//      }
+//      bb4: {
+//          _5 = const loop_forever() -> bb5;
+//      }
+//      bb5: {
+//          StorageDead(_6);
+//          StorageDead(_5);
+//          StorageLive(_7);
+//          unreachable;
+//      }
+//  }
+// END rustc.main.UnreachablePropagation.after.mir