JumpThreading: Re-enable and fix Not ops on non-booleans

clubby789 · clubby789 · commit 9f35fe47c75c · 2025-04-13T20:29:49.000Z
diff --git a/compiler/rustc_mir_transform/src/jump_threading.rs b/compiler/rustc_mir_transform/src/jump_threading.rs
@@ -150,14 +150,6 @@ impl Condition {
     fn matches(&self, value: ScalarInt) -> bool {
         (self.value == value) == (self.polarity == Polarity::Eq)
     }
-
-    fn inv(mut self) -> Self {
-        self.polarity = match self.polarity {
-            Polarity::Eq => Polarity::Ne,
-            Polarity::Ne => Polarity::Eq,
-        };
-        self
-    }
 }
 
 #[derive(Copy, Clone, Debug)]
@@ -495,19 +487,20 @@ impl<'a, 'tcx> TOFinder<'a, 'tcx> {
                     }
                 }
             }
-            // Transfer the conditions on the copy rhs, after inversing polarity.
+            // Transfer the conditions on the copy rhs, after inverting the value of the condition.
             Rvalue::UnaryOp(UnOp::Not, Operand::Move(place) | Operand::Copy(place)) => {
-                if !place.ty(self.body, self.tcx).ty.is_bool() {
-                    // Constructing the conditions by inverting the polarity
-                    // of equality is only correct for bools. That is to say,
-                    // `!a == b` is not `a != b` for integers greater than 1 bit.
-                    return;
-                }
+                let layout = self.ecx.layout_of(place.ty(self.body, self.tcx).ty).unwrap();
                 let Some(conditions) = state.try_get_idx(lhs, &self.map) else { return };
                 let Some(place) = self.map.find(place.as_ref()) else { return };
-                // FIXME: I think This could be generalized to not bool if we
-                // actually perform a logical not on the condition's value.
-                let conds = conditions.map(self.arena, Condition::inv);
+                let conds = conditions.map(self.arena, |mut cond| {
+                    cond.value = self
+                        .ecx
+                        .unary_op(UnOp::Not, &ImmTy::from_scalar_int(cond.value, layout))
+                        .unwrap()
+                        .to_scalar_int()
+                        .unwrap();
+                    cond
+                });
                 state.insert_value_idx(place, conds, &self.map);
             }
             // We expect `lhs ?= A`. We found `lhs = Eq(rhs, B)`.
diff --git a/tests/mir-opt/jump_threading.bitwise_not.JumpThreading.panic-abort.diff b/tests/mir-opt/jump_threading.bitwise_not.JumpThreading.panic-abort.diff
@@ -3,7 +3,7 @@
   
   fn bitwise_not() -> i32 {
       let mut _0: i32;
-      let mut _1: i32;
+      let _1: i32;
       let mut _2: bool;
       let mut _3: i32;
       let mut _4: i32;
@@ -13,7 +13,6 @@
   
       bb0: {
           StorageLive(_1);
-          _1 = const 0_i32;
           _1 = const 1_i32;
           StorageLive(_2);
           StorageLive(_3);
@@ -22,7 +21,8 @@
           _3 = Not(move _4);
           StorageDead(_4);
           _2 = Eq(move _3, const 0_i32);
-          switchInt(move _2) -> [0: bb2, otherwise: bb1];
+-         switchInt(move _2) -> [0: bb2, otherwise: bb1];
++         goto -> bb2;
       }
   
       bb1: {
diff --git a/tests/mir-opt/jump_threading.bitwise_not.JumpThreading.panic-unwind.diff b/tests/mir-opt/jump_threading.bitwise_not.JumpThreading.panic-unwind.diff
@@ -3,7 +3,7 @@
   
   fn bitwise_not() -> i32 {
       let mut _0: i32;
-      let mut _1: i32;
+      let _1: i32;
       let mut _2: bool;
       let mut _3: i32;
       let mut _4: i32;
@@ -13,7 +13,6 @@
   
       bb0: {
           StorageLive(_1);
-          _1 = const 0_i32;
           _1 = const 1_i32;
           StorageLive(_2);
           StorageLive(_3);
@@ -22,7 +21,8 @@
           _3 = Not(move _4);
           StorageDead(_4);
           _2 = Eq(move _3, const 0_i32);
-          switchInt(move _2) -> [0: bb2, otherwise: bb1];
+-         switchInt(move _2) -> [0: bb2, otherwise: bb1];
++         goto -> bb2;
       }
   
       bb1: {
diff --git a/tests/mir-opt/jump_threading.logical_not.JumpThreading.panic-abort.diff b/tests/mir-opt/jump_threading.logical_not.JumpThreading.panic-abort.diff
@@ -0,0 +1,46 @@
+- // MIR for `logical_not` before JumpThreading
++ // MIR for `logical_not` after JumpThreading
+  
+  fn logical_not() -> i32 {
+      let mut _0: i32;
+      let _1: bool;
+      let mut _2: bool;
+      let mut _3: bool;
+      let mut _4: bool;
+      scope 1 {
+          debug a => _1;
+      }
+  
+      bb0: {
+          StorageLive(_1);
+          _1 = const false;
+          StorageLive(_2);
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = copy _1;
+          _3 = Not(move _4);
+          StorageDead(_4);
+          _2 = Eq(move _3, const true);
+-         switchInt(move _2) -> [0: bb2, otherwise: bb1];
++         goto -> bb1;
+      }
+  
+      bb1: {
+          StorageDead(_3);
+          _0 = const 1_i32;
+          goto -> bb3;
+      }
+  
+      bb2: {
+          StorageDead(_3);
+          _0 = const 0_i32;
+          goto -> bb3;
+      }
+  
+      bb3: {
+          StorageDead(_2);
+          StorageDead(_1);
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/jump_threading.logical_not.JumpThreading.panic-unwind.diff b/tests/mir-opt/jump_threading.logical_not.JumpThreading.panic-unwind.diff
@@ -0,0 +1,46 @@
+- // MIR for `logical_not` before JumpThreading
++ // MIR for `logical_not` after JumpThreading
+  
+  fn logical_not() -> i32 {
+      let mut _0: i32;
+      let _1: bool;
+      let mut _2: bool;
+      let mut _3: bool;
+      let mut _4: bool;
+      scope 1 {
+          debug a => _1;
+      }
+  
+      bb0: {
+          StorageLive(_1);
+          _1 = const false;
+          StorageLive(_2);
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = copy _1;
+          _3 = Not(move _4);
+          StorageDead(_4);
+          _2 = Eq(move _3, const true);
+-         switchInt(move _2) -> [0: bb2, otherwise: bb1];
++         goto -> bb1;
+      }
+  
+      bb1: {
+          StorageDead(_3);
+          _0 = const 1_i32;
+          goto -> bb3;
+      }
+  
+      bb2: {
+          StorageDead(_3);
+          _0 = const 0_i32;
+          goto -> bb3;
+      }
+  
+      bb3: {
+          StorageDead(_2);
+          StorageDead(_1);
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/jump_threading.rs b/tests/mir-opt/jump_threading.rs
@@ -532,14 +532,19 @@ fn floats() -> u32 {
 
 pub fn bitwise_not() -> i32 {
     // CHECK-LABEL: fn bitwise_not(
-    // CHECK: switchInt(
 
     // Test for #131195, which was optimizing `!a == b` into `a != b`.
-    let mut a: i32 = 0;
-    a = 1;
+    let a = 1;
     if !a == 0 { 1 } else { 0 }
 }
 
+pub fn logical_not() -> i32 {
+    // CHECK-LABEL: fn logical_not(
+
+    let a = false;
+    if !a == true { 1 } else { 0 }
+}
+
 fn main() {
     // CHECK-LABEL: fn main(
     too_complex(Ok(0));
@@ -555,6 +560,8 @@ fn main() {
     aggregate(7);
     assume(7, false);
     floats();
+    bitwise_not();
+    logical_not();
 }
 
 // EMIT_MIR jump_threading.too_complex.JumpThreading.diff
@@ -572,3 +579,4 @@ fn main() {
 // EMIT_MIR jump_threading.aggregate_copy.JumpThreading.diff
 // EMIT_MIR jump_threading.floats.JumpThreading.diff
 // EMIT_MIR jump_threading.bitwise_not.JumpThreading.diff
+// EMIT_MIR jump_threading.logical_not.JumpThreading.diff
