Simplify assume of a constant.

cjgillot · cjgillot · commit 0b13e636f54e · 2023-10-31T11:44:23.000Z
diff --git a/compiler/rustc_mir_transform/src/simplify_branches.rs b/compiler/rustc_mir_transform/src/simplify_branches.rs
@@ -16,8 +16,25 @@ impl<'tcx> MirPass<'tcx> for SimplifyConstCondition {
     }
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        trace!("Running SimplifyConstCondition on {:?}", body.source);
         let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
-        for block in body.basic_blocks_mut() {
+        'blocks: for block in body.basic_blocks_mut() {
+            for stmt in block.statements.iter_mut() {
+                if let StatementKind::Intrinsic(box ref intrinsic) = stmt.kind
+                    && let NonDivergingIntrinsic::Assume(discr) = intrinsic
+                    && let Operand::Constant(ref c) = discr
+                    && let Some(constant) = c.const_.try_eval_bool(tcx, param_env)
+                {
+                    if constant {
+                        stmt.make_nop();
+                    } else {
+                        block.statements.clear();
+                        block.terminator_mut().kind = TerminatorKind::Unreachable;
+                        continue 'blocks;
+                    }
+                }
+            }
+
             let terminator = block.terminator_mut();
             terminator.kind = match terminator.kind {
                 TerminatorKind::SwitchInt {
diff --git a/tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.Inline.panic-abort.diff b/tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.Inline.panic-abort.diff
@@ -9,14 +9,15 @@
 +         debug self => _2;
 +         let mut _3: &std::option::Option<T>;
 +         let mut _4: isize;
++         let mut _5: bool;
++         let mut _6: bool;
 +         scope 2 {
 +             debug val => _0;
 +         }
 +         scope 3 {
 +             scope 5 (inlined unreachable_unchecked) {
 +                 scope 6 {
 +                     scope 7 (inlined unreachable_unchecked::runtime) {
-+                         let _5: !;
 +                     }
 +                 }
 +             }
@@ -30,24 +31,20 @@
           StorageLive(_2);
           _2 = move _1;
 -         _0 = Option::<T>::unwrap_unchecked(move _2) -> [return: bb1, unwind unreachable];
+-     }
+- 
+-     bb1: {
 +         StorageLive(_3);
 +         StorageLive(_4);
 +         StorageLive(_5);
++         StorageLive(_6);
 +         _4 = discriminant(_2);
-+         switchInt(move _4) -> [0: bb1, 1: bb3, otherwise: bb2];
-      }
-  
-      bb1: {
-+         assume(const false);
-+         _5 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
-+     }
-+ 
-+     bb2: {
-+         unreachable;
-+     }
-+ 
-+     bb3: {
++         _5 = Ne(_4, const 0_isize);
++         assume(move _5);
++         _6 = Eq(_4, const 1_isize);
++         assume(move _6);
 +         _0 = move ((_2 as Some).0: T);
++         StorageDead(_6);
 +         StorageDead(_5);
 +         StorageDead(_4);
 +         StorageDead(_3);
diff --git a/tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.Inline.panic-unwind.diff b/tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.Inline.panic-unwind.diff
@@ -9,14 +9,15 @@
 +         debug self => _2;
 +         let mut _3: &std::option::Option<T>;
 +         let mut _4: isize;
++         let mut _5: bool;
++         let mut _6: bool;
 +         scope 2 {
 +             debug val => _0;
 +         }
 +         scope 3 {
 +             scope 5 (inlined unreachable_unchecked) {
 +                 scope 6 {
 +                     scope 7 (inlined unreachable_unchecked::runtime) {
-+                         let _5: !;
 +                     }
 +                 }
 +             }
@@ -30,33 +31,29 @@
           StorageLive(_2);
           _2 = move _1;
 -         _0 = Option::<T>::unwrap_unchecked(move _2) -> [return: bb1, unwind: bb2];
+-     }
+- 
+-     bb1: {
 +         StorageLive(_3);
 +         StorageLive(_4);
 +         StorageLive(_5);
++         StorageLive(_6);
 +         _4 = discriminant(_2);
-+         switchInt(move _4) -> [0: bb1, 1: bb3, otherwise: bb2];
-      }
-  
-      bb1: {
--         StorageDead(_2);
--         return;
-+         assume(const false);
-+         _5 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
-      }
-  
--     bb2 (cleanup): {
--         resume;
-+     bb2: {
-+         unreachable;
-+     }
-+ 
-+     bb3: {
++         _5 = Ne(_4, const 0_isize);
++         assume(move _5);
++         _6 = Eq(_4, const 1_isize);
++         assume(move _6);
 +         _0 = move ((_2 as Some).0: T);
++         StorageDead(_6);
 +         StorageDead(_5);
 +         StorageDead(_4);
 +         StorageDead(_3);
-+         StorageDead(_2);
-+         return;
+          StorageDead(_2);
+          return;
+-     }
+- 
+-     bb2 (cleanup): {
+-         resume;
       }
   }
   
diff --git a/tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.PreCodegen.after.panic-abort.mir b/tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.PreCodegen.after.panic-abort.mir
@@ -6,44 +6,40 @@ fn unwrap_unchecked(_1: Option<T>) -> T {
     scope 1 (inlined #[track_caller] Option::<T>::unwrap_unchecked) {
         debug self => _1;
         let mut _2: isize;
-        let mut _4: &std::option::Option<T>;
+        let mut _3: bool;
+        let mut _4: bool;
+        let mut _5: &std::option::Option<T>;
         scope 2 {
             debug val => _0;
         }
         scope 3 {
             scope 5 (inlined unreachable_unchecked) {
                 scope 6 {
                     scope 7 (inlined unreachable_unchecked::runtime) {
-                        let _3: !;
                     }
                 }
             }
         }
         scope 4 (inlined Option::<T>::is_some) {
-            debug self => _4;
+            debug self => _5;
         }
     }
 
     bb0: {
-        StorageLive(_4);
+        StorageLive(_5);
         StorageLive(_2);
+        StorageLive(_3);
+        StorageLive(_4);
         _2 = discriminant(_1);
-        switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb3];
-    }
-
-    bb1: {
-        assume(const false);
-        _3 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
-    }
-
-    bb2: {
+        _3 = Ne(_2, const 0_isize);
+        assume(move _3);
+        _4 = Eq(_2, const 1_isize);
+        assume(move _4);
         _0 = move ((_1 as Some).0: T);
-        StorageDead(_2);
         StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_2);
+        StorageDead(_5);
         return;
     }
-
-    bb3: {
-        unreachable;
-    }
 }
diff --git a/tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.PreCodegen.after.panic-unwind.mir b/tests/mir-opt/inline/unwrap_unchecked.unwrap_unchecked.PreCodegen.after.panic-unwind.mir
@@ -6,44 +6,40 @@ fn unwrap_unchecked(_1: Option<T>) -> T {
     scope 1 (inlined #[track_caller] Option::<T>::unwrap_unchecked) {
         debug self => _1;
         let mut _2: isize;
-        let mut _4: &std::option::Option<T>;
+        let mut _3: bool;
+        let mut _4: bool;
+        let mut _5: &std::option::Option<T>;
         scope 2 {
             debug val => _0;
         }
         scope 3 {
             scope 5 (inlined unreachable_unchecked) {
                 scope 6 {
                     scope 7 (inlined unreachable_unchecked::runtime) {
-                        let _3: !;
                     }
                 }
             }
         }
         scope 4 (inlined Option::<T>::is_some) {
-            debug self => _4;
+            debug self => _5;
         }
     }
 
     bb0: {
-        StorageLive(_4);
+        StorageLive(_5);
         StorageLive(_2);
+        StorageLive(_3);
+        StorageLive(_4);
         _2 = discriminant(_1);
-        switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb3];
-    }
-
-    bb1: {
-        assume(const false);
-        _3 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
-    }
-
-    bb2: {
+        _3 = Ne(_2, const 0_isize);
+        assume(move _3);
+        _4 = Eq(_2, const 1_isize);
+        assume(move _4);
         _0 = move ((_1 as Some).0: T);
-        StorageDead(_2);
         StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_2);
+        StorageDead(_5);
         return;
     }
-
-    bb3: {
-        unreachable;
-    }
 }
diff --git a/tests/mir-opt/pre-codegen/duplicate_switch_targets.ub_if_b.PreCodegen.after.mir b/tests/mir-opt/pre-codegen/duplicate_switch_targets.ub_if_b.PreCodegen.after.mir
@@ -4,30 +4,22 @@ fn ub_if_b(_1: Thing) -> Thing {
     debug t => _1;
     let mut _0: Thing;
     let mut _2: isize;
+    let mut _3: bool;
+    let mut _4: bool;
     scope 1 (inlined unreachable_unchecked) {
         scope 2 {
             scope 3 (inlined unreachable_unchecked::runtime) {
-                let _3: !;
             }
         }
     }
 
     bb0: {
         _2 = discriminant(_1);
-        switchInt(move _2) -> [0: bb1, 1: bb2, otherwise: bb3];
-    }
-
-    bb1: {
+        _3 = Ne(_2, const 1_isize);
+        assume(move _3);
+        _4 = Eq(_2, const 0_isize);
+        assume(move _4);
         _0 = move _1;
         return;
     }
-
-    bb2: {
-        assume(const false);
-        _3 = core::panicking::panic_nounwind(const "unsafe precondition(s) violated: hint::unreachable_unchecked must never be reached") -> unwind unreachable;
-    }
-
-    bb3: {
-        unreachable;
-    }
 }
