mir-opt: Simplify trivial constants in SimplifyConstCondition

Author: dianqk

compiler/rustc_mir_transform/src/lib.rs

@@ -189,6 +189,7 @@ declare_passes! {
             Final
         };
     mod simplify_branches : SimplifyConstCondition {
+        AfterInstSimplify,
         AfterConstProp,
         Final
     };
@@ -708,6 +709,7 @@ 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,
+            &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",
         }
@@ -24,19 +28,39 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
         let mut patch = MirPatch::new(body);
 
         'blocks: for (bb, block) in body.basic_blocks.iter_enumerated() {
+            let mut pre_local_const: Option<(Local, &'_ ConstOperand<'_>)> = None;
+
             for (statement_index, stmt) in block.statements.iter().enumerate() {
+                let has_local_const = pre_local_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(constant) = c.const_.try_eval_bool(tcx, typing_env)
                 {
+                    let c = if let Operand::Constant(c) = discr {
+                        c
+                    } else if let Some((local, c)) = has_local_const
+                        && let Some(assume_local) = discr.place().and_then(|p| p.as_local())
+                        && local == assume_local
+                    {
+                        c
+                    } else {
+                        continue;
+                    };
+                    let Some(constant) = c.const_.try_eval_bool(tcx, typing_env) else {
+                        continue;
+                    };
                     if constant {
                         patch.nop_statement(Location { block: bb, statement_index });
                     } else {
                         patch.patch_terminator(bb, TerminatorKind::Unreachable);
                         continue 'blocks;
                     }
+                } else if let StatementKind::Assign(box (ref lhs, ref rvalue)) = stmt.kind
+                    && let Some(local) = lhs.as_local()
+                    && let Rvalue::Use(Operand::Constant(c)) = rvalue
+                    && c.const_.ty().is_bool()
+                {
+                    pre_local_const = Some((local, c));
                 }
             }
 

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/pre-codegen/two_unwrap_unchecked.rs

@@ -0,0 +1,12 @@
+// skip-filecheck
+//@ compile-flags: -O
+
+#![crate_type = "lib"]
+
+// EMIT_MIR two_unwrap_unchecked.two_unwrap_unchecked.GVN.diff
+// EMIT_MIR two_unwrap_unchecked.two_unwrap_unchecked.PreCodegen.after.mir
+pub fn two_unwrap_unchecked(v: &Option<i32>) -> i32 {
+    let v1 = unsafe { v.unwrap_unchecked() };
+    let v2 = unsafe { v.unwrap_unchecked() };
+    v1 + v2
+}

tests/mir-opt/pre-codegen/two_unwrap_unchecked.two_unwrap_unchecked.GVN.diff

@@ -0,0 +1,96 @@
+- // MIR for `two_unwrap_unchecked` before GVN
++ // MIR for `two_unwrap_unchecked` after GVN
+  
+  fn two_unwrap_unchecked(_1: &Option<i32>) -> i32 {
+      debug v => _1;
+      let mut _0: i32;
+      let _2: i32;
+      let mut _3: std::option::Option<i32>;
+      let mut _5: std::option::Option<i32>;
+      let mut _6: i32;
+      let mut _7: i32;
+      scope 1 {
+          debug v1 => _2;
+          let _4: i32;
+          scope 2 {
+              debug v2 => _4;
+          }
+          scope 8 (inlined #[track_caller] Option::<i32>::unwrap_unchecked) {
+              let mut _9: isize;
+              scope 9 {
+              }
+              scope 10 (inlined #[track_caller] unreachable_unchecked) {
+                  scope 11 (inlined core::ub_checks::check_language_ub) {
+                      scope 12 (inlined core::ub_checks::check_language_ub::runtime) {
+                      }
+                  }
+              }
+          }
+      }
+      scope 3 (inlined #[track_caller] Option::<i32>::unwrap_unchecked) {
+          let mut _8: isize;
+          scope 4 {
+          }
+          scope 5 (inlined #[track_caller] unreachable_unchecked) {
+              scope 6 (inlined core::ub_checks::check_language_ub) {
+                  scope 7 (inlined core::ub_checks::check_language_ub::runtime) {
+                  }
+              }
+          }
+      }
+  
+      bb0: {
+-         StorageLive(_2);
++         nop;
+          StorageLive(_3);
+          _3 = copy (*_1);
+          StorageLive(_8);
+          _8 = discriminant(_3);
+          switchInt(move _8) -> [0: bb2, 1: bb3, otherwise: bb1];
+      }
+  
+      bb1: {
+          unreachable;
+      }
+  
+      bb2: {
+          unreachable;
+      }
+  
+      bb3: {
+          _2 = move ((_3 as Some).0: i32);
+          StorageDead(_8);
+          StorageDead(_3);
+-         StorageLive(_4);
++         nop;
+          StorageLive(_5);
+          _5 = copy (*_1);
+          StorageLive(_9);
+          _9 = discriminant(_5);
+          switchInt(move _9) -> [0: bb4, 1: bb5, otherwise: bb1];
+      }
+  
+      bb4: {
+          unreachable;
+      }
+  
+      bb5: {
+          _4 = move ((_5 as Some).0: i32);
+          StorageDead(_9);
+          StorageDead(_5);
+          StorageLive(_6);
+          _6 = copy _2;
+          StorageLive(_7);
+          _7 = copy _4;
+-         _0 = Add(move _6, move _7);
++         _0 = Add(copy _2, copy _4);
+          StorageDead(_7);
+          StorageDead(_6);
+-         StorageDead(_4);
+-         StorageDead(_2);
++         nop;
++         nop;
+          return;
+      }
+  }
+  

tests/mir-opt/pre-codegen/two_unwrap_unchecked.two_unwrap_unchecked.PreCodegen.after.mir

@@ -0,0 +1,69 @@
+// MIR for `two_unwrap_unchecked` after PreCodegen
+
+fn two_unwrap_unchecked(_1: &Option<i32>) -> i32 {
+    debug v => _1;
+    let mut _0: i32;
+    let mut _2: std::option::Option<i32>;
+    let _4: i32;
+    let mut _5: std::option::Option<i32>;
+    scope 1 {
+        debug v1 => _4;
+        let _7: i32;
+        scope 2 {
+            debug v2 => _7;
+        }
+        scope 8 (inlined #[track_caller] Option::<i32>::unwrap_unchecked) {
+            let mut _6: isize;
+            scope 9 {
+            }
+            scope 10 (inlined #[track_caller] unreachable_unchecked) {
+                scope 11 (inlined core::ub_checks::check_language_ub) {
+                    scope 12 (inlined core::ub_checks::check_language_ub::runtime) {
+                    }
+                }
+            }
+        }
+    }
+    scope 3 (inlined #[track_caller] Option::<i32>::unwrap_unchecked) {
+        let mut _3: isize;
+        scope 4 {
+        }
+        scope 5 (inlined #[track_caller] unreachable_unchecked) {
+            scope 6 (inlined core::ub_checks::check_language_ub) {
+                scope 7 (inlined core::ub_checks::check_language_ub::runtime) {
+                }
+            }
+        }
+    }
+
+    bb0: {
+        StorageLive(_2);
+        _2 = copy (*_1);
+        StorageLive(_3);
+        _3 = discriminant(_2);
+        switchInt(move _3) -> [0: bb3, 1: bb1, otherwise: bb3];
+    }
+
+    bb1: {
+        _4 = move ((_2 as Some).0: i32);
+        StorageDead(_3);
+        StorageDead(_2);
+        StorageLive(_5);
+        _5 = copy (*_1);
+        StorageLive(_6);
+        _6 = discriminant(_5);
+        switchInt(move _6) -> [0: bb3, 1: bb2, otherwise: bb3];
+    }
+
+    bb2: {
+        _7 = move ((_5 as Some).0: i32);
+        StorageDead(_6);
+        StorageDead(_5);
+        _0 = Add(copy _4, copy _7);
+        return;
+    }
+
+    bb3: {
+        unreachable;
+    }
+}