Auto merge of rust-lang#147654 - dianqk:simplify-const-condition, r=cjgillot

Simplify trivial constants in SimplifyConstCondition

After `InstSimplify-after-simplifycfg` with `-Zub_checks=false`, there are many of the following patterns.

```
_13 = const false;
assume(copy _13);
_12 = unreachable_unchecked::precondition_check() -> [return: bb1, unwind unreachable];
```

Simplifying them to unreachable early should make CFG simpler.

Author: bors

compiler/rustc_middle/src/mir/terminator.rs

@@ -696,6 +696,28 @@ impl<'tcx> TerminatorKind<'tcx> {
             _ => None,
         }
     }
+
+    /// Returns true if the terminator can write to memory.
+    pub fn can_write_to_memory(&self) -> bool {
+        match self {
+            TerminatorKind::Goto { .. }
+            | TerminatorKind::SwitchInt { .. }
+            | TerminatorKind::UnwindResume
+            | TerminatorKind::UnwindTerminate(_)
+            | TerminatorKind::Return
+            | TerminatorKind::Assert { .. }
+            | TerminatorKind::CoroutineDrop
+            | TerminatorKind::FalseEdge { .. }
+            | TerminatorKind::FalseUnwind { .. }
+            | TerminatorKind::Unreachable => false,
+            TerminatorKind::Call { .. }
+            | TerminatorKind::Drop { .. }
+            | TerminatorKind::TailCall { .. }
+            // Yield writes to the resume_arg place.
+            | TerminatorKind::Yield { .. }
+            | TerminatorKind::InlineAsm { .. } => true,
+        }
+    }
 }
 
 #[derive(Copy, Clone, Debug)]

compiler/rustc_mir_transform/src/gvn.rs

@@ -1926,13 +1926,8 @@ impl<'tcx> MutVisitor<'tcx> for VnState<'_, '_, 'tcx> {
                 self.assign(local, opaque);
             }
         }
-        // Function calls and ASM may invalidate (nested) derefs. We must handle them carefully.
-        // Currently, only preserving derefs for trivial terminators like SwitchInt and Goto.
-        let safe_to_preserve_derefs = matches!(
-            terminator.kind,
-            TerminatorKind::SwitchInt { .. } | TerminatorKind::Goto { .. }
-        );
-        if !safe_to_preserve_derefs {
+        // Terminators that can write to memory may invalidate (nested) derefs.
+        if terminator.kind.can_write_to_memory() {
             self.invalidate_derefs();
         }
         self.super_terminator(terminator, location);

compiler/rustc_mir_transform/src/lib.rs

@@ -189,6 +189,7 @@ declare_passes! {
             Final
         };
     mod simplify_branches : SimplifyConstCondition {
+        AfterInstSimplify,
         AfterConstProp,
         Final
     };
@@ -708,6 +709,15 @@ pub(crate) fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'
             // optimizations. This invalidates CFG caches, so avoid putting between
             // `ReferencePropagation` and `GVN` which both use the dominator tree.
             &instsimplify::InstSimplify::AfterSimplifyCfg,
+            // After `InstSimplify-after-simplifycfg` with `-Zub_checks=false`, simplify
+            // ```
+            // _13 = const false;
+            // assume(copy _13);
+            // Call(precondition_check);
+            // ```
+            // to unreachable to eliminate the call to help later passes.
+            // This invalidates CFG caches also.
+            &o1(simplify_branches::SimplifyConstCondition::AfterInstSimplify),
             &ref_prop::ReferencePropagation,
             &sroa::ScalarReplacementOfAggregates,
             &simplify::SimplifyLocals::BeforeConstProp,

compiler/rustc_mir_transform/src/simplify_branches.rs

@@ -5,6 +5,7 @@ use tracing::trace;
 use crate::patch::MirPatch;
 
 pub(super) enum SimplifyConstCondition {
+    AfterInstSimplify,
     AfterConstProp,
     Final,
 }
@@ -13,6 +14,9 @@ pub(super) enum SimplifyConstCondition {
 impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
     fn name(&self) -> &'static str {
         match self {
+            SimplifyConstCondition::AfterInstSimplify => {
+                "SimplifyConstCondition-after-inst-simplify"
+            }
             SimplifyConstCondition::AfterConstProp => "SimplifyConstCondition-after-const-prop",
             SimplifyConstCondition::Final => "SimplifyConstCondition-final",
         }
@@ -23,12 +27,33 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
         let typing_env = body.typing_env(tcx);
         let mut patch = MirPatch::new(body);
 
+        fn try_get_const<'tcx, 'a>(
+            operand: &'a Operand<'tcx>,
+            has_place_const: Option<(Place<'tcx>, &'a ConstOperand<'tcx>)>,
+        ) -> Option<&'a ConstOperand<'tcx>> {
+            match operand {
+                Operand::Constant(const_operand) => Some(const_operand),
+                // `has_place_const` must be the LHS of the previous statement.
+                // Soundness: There is nothing can modify the place, as there are no statements between the two statements.
+                Operand::Copy(place) | Operand::Move(place)
+                    if let Some((place_const, const_operand)) = has_place_const
+                        && place_const == *place =>
+                {
+                    Some(const_operand)
+                }
+                Operand::Copy(_) | Operand::Move(_) => None,
+            }
+        }
+
         'blocks: for (bb, block) in body.basic_blocks.iter_enumerated() {
+            let mut pre_place_const: Option<(Place<'tcx>, &ConstOperand<'tcx>)> = None;
+
             for (statement_index, stmt) in block.statements.iter().enumerate() {
+                let has_place_const = pre_place_const.take();
                 // Simplify `assume` of a known value: either a NOP or unreachable.
                 if let StatementKind::Intrinsic(box ref intrinsic) = stmt.kind
                     && let NonDivergingIntrinsic::Assume(discr) = intrinsic
-                    && let Operand::Constant(c) = discr
+                    && let Some(c) = try_get_const(discr, has_place_const)
                     && let Some(constant) = c.const_.try_eval_bool(tcx, typing_env)
                 {
                     if constant {
@@ -37,28 +62,29 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
                         patch.patch_terminator(bb, TerminatorKind::Unreachable);
                         continue 'blocks;
                     }
+                } else if let StatementKind::Assign(box (lhs, ref rvalue)) = stmt.kind
+                    && let Rvalue::Use(Operand::Constant(c)) = rvalue
+                {
+                    pre_place_const = Some((lhs, c));
                 }
             }
 
             let terminator = block.terminator();
             let terminator = match terminator.kind {
-                TerminatorKind::SwitchInt {
-                    discr: Operand::Constant(ref c), ref targets, ..
-                } => {
-                    let constant = c.const_.try_eval_bits(tcx, typing_env);
-                    if let Some(constant) = constant {
-                        let target = targets.target_for_value(constant);
-                        TerminatorKind::Goto { target }
-                    } else {
-                        continue;
-                    }
+                TerminatorKind::SwitchInt { ref discr, ref targets, .. }
+                    if let Some(c) = try_get_const(discr, pre_place_const.take())
+                        && let Some(constant) = c.const_.try_eval_bits(tcx, typing_env) =>
+                {
+                    let target = targets.target_for_value(constant);
+                    TerminatorKind::Goto { target }
+                }
+                TerminatorKind::Assert { target, ref cond, expected, .. }
+                    if let Some(c) = try_get_const(&cond, pre_place_const.take())
+                        && let Some(constant) = c.const_.try_eval_bool(tcx, typing_env)
+                        && constant == expected =>
+                {
+                    TerminatorKind::Goto { target }
                 }
-                TerminatorKind::Assert {
-                    target, cond: Operand::Constant(ref c), expected, ..
-                } => match c.const_.try_eval_bool(tcx, typing_env) {
-                    Some(v) if v == expected => TerminatorKind::Goto { target },
-                    _ => continue,
-                },
                 _ => continue,
             };
             patch.patch_terminator(bb, terminator);

tests/mir-opt/const_prop/control_flow_simplification.hello.GVN.panic-abort.diff

@@ -3,23 +3,17 @@
 
   fn hello() -> () {
       let mut _0: ();
-      let mut _1: bool;
-      let mut _2: !;
+      let mut _1: !;
 
       bb0: {
-          StorageLive(_1);
--         _1 = const <bool as NeedsDrop>::NEEDS;
--         switchInt(move _1) -> [0: bb2, otherwise: bb1];
-+         _1 = const false;
-+         switchInt(const false) -> [0: bb2, otherwise: bb1];
+          goto -> bb2;
       }
 
       bb1: {
-          _2 = begin_panic::<&str>(const "explicit panic") -> unwind unreachable;
+          _1 = begin_panic::<&str>(const "explicit panic") -> unwind unreachable;
       }
 
       bb2: {
-          StorageDead(_1);
           return;
       }
   }

tests/mir-opt/const_prop/control_flow_simplification.hello.GVN.panic-unwind.diff

@@ -3,23 +3,17 @@
 
   fn hello() -> () {
       let mut _0: ();
-      let mut _1: bool;
-      let mut _2: !;
+      let mut _1: !;
 
       bb0: {
-          StorageLive(_1);
--         _1 = const <bool as NeedsDrop>::NEEDS;
--         switchInt(move _1) -> [0: bb2, otherwise: bb1];
-+         _1 = const false;
-+         switchInt(const false) -> [0: bb2, otherwise: bb1];
+          goto -> bb2;
       }
 
       bb1: {
-          _2 = begin_panic::<&str>(const "explicit panic") -> unwind continue;
+          _1 = begin_panic::<&str>(const "explicit panic") -> unwind continue;
       }
 
       bb2: {
-          StorageDead(_1);
           return;
       }
   }

tests/mir-opt/const_prop/control_flow_simplification.rs

@@ -1,4 +1,3 @@
-// skip-filecheck
 // EMIT_MIR_FOR_EACH_PANIC_STRATEGY
 //@ test-mir-pass: GVN
 //@ compile-flags: -Zmir-opt-level=1
@@ -12,6 +11,9 @@ impl<This> NeedsDrop for This {}
 // EMIT_MIR control_flow_simplification.hello.GVN.diff
 // EMIT_MIR control_flow_simplification.hello.PreCodegen.before.mir
 fn hello<T>() {
+    // CHECK-LABEL: fn hello(
+    // CHECK: bb0:
+    // CHECK-NEXT: return;
     if <bool>::NEEDS {
         panic!()
     }

tests/mir-opt/const_prop/trivial_const.rs

@@ -0,0 +1,15 @@
+//@ test-mir-pass: SimplifyConstCondition-after-inst-simplify
+//@ compile-flags: -Zmir-enable-passes=+InstSimplify-after-simplifycfg -Zub_checks=false -Zinline-mir
+
+#![crate_type = "lib"]
+
+// EMIT_MIR trivial_const.unwrap_unchecked.SimplifyConstCondition-after-inst-simplify.diff
+pub fn unwrap_unchecked(v: &Option<i32>) -> i32 {
+    // CHECK-LABEL: fn unwrap_unchecked(
+    // CHECK: bb0: {
+    // CHECK: switchInt({{.*}}) -> [0: [[AssumeFalseBB:bb.*]], 1:
+    // CHECK: [[AssumeFalseBB]]: {
+    // CHECK-NEXT: unreachable;
+    let v = unsafe { v.unwrap_unchecked() };
+    v
+}

tests/mir-opt/const_prop/trivial_const.unwrap_unchecked.SimplifyConstCondition-after-inst-simplify.diff

@@ -0,0 +1,58 @@
+- // MIR for `unwrap_unchecked` before SimplifyConstCondition-after-inst-simplify
++ // MIR for `unwrap_unchecked` after SimplifyConstCondition-after-inst-simplify
+  
+  fn unwrap_unchecked(_1: &Option<i32>) -> i32 {
+      debug v => _1;
+      let mut _0: i32;
+      let _2: i32;
+      let mut _3: std::option::Option<i32>;
+      scope 1 {
+          debug v => _2;
+      }
+      scope 2 (inlined #[track_caller] Option::<i32>::unwrap_unchecked) {
+          let mut _4: isize;
+          scope 3 {
+          }
+          scope 4 (inlined #[track_caller] unreachable_unchecked) {
+              let _5: ();
+              scope 5 (inlined core::ub_checks::check_language_ub) {
+                  let mut _6: bool;
+                  scope 6 (inlined core::ub_checks::check_language_ub::runtime) {
+                  }
+              }
+          }
+      }
+  
+      bb0: {
+          StorageLive(_2);
+          StorageLive(_3);
+          _3 = copy (*_1);
+          StorageLive(_4);
+          StorageLive(_5);
+          _4 = discriminant(_3);
+          switchInt(move _4) -> [0: bb2, 1: bb3, otherwise: bb1];
+      }
+  
+      bb1: {
+          unreachable;
+      }
+  
+      bb2: {
+-         StorageLive(_6);
+-         _6 = const false;
+-         assume(copy _6);
+-         _5 = unreachable_unchecked::precondition_check() -> [return: bb1, unwind unreachable];
++         unreachable;
+      }
+  
+      bb3: {
+          _2 = move ((_3 as Some).0: i32);
+          StorageDead(_5);
+          StorageDead(_4);
+          StorageDead(_3);
+          _0 = copy _2;
+          StorageDead(_2);
+          return;
+      }
+  }
+  

tests/mir-opt/dont_reset_cast_kind_without_updating_operand.rs

@@ -1,10 +1,14 @@
-// skip-filecheck
-//@ compile-flags: -Zmir-enable-passes=+Inline,+GVN --crate-type lib
+//@ test-mir-pass: GVN
+//@ compile-flags: -Zinline-mir --crate-type lib
 // EMIT_MIR_FOR_EACH_BIT_WIDTH
 // EMIT_MIR_FOR_EACH_PANIC_STRATEGY
 // EMIT_MIR dont_reset_cast_kind_without_updating_operand.test.GVN.diff
 
 fn test() {
+    // CHECK-LABEL: fn test(
+    // CHECK: debug slf => [[SLF:_.*]];
+    // CHECK: debug _x => [[X:_.*]];
+    // CHECK: [[X]] = copy [[SLF]] as *mut () (PtrToPtr);
     let vp_ctx: &Box<()> = &Box::new(());
     let slf: *const () = &raw const **vp_ctx;
     let bytes = std::ptr::slice_from_raw_parts(slf, 1);