From fd35e820bbb0385c61246e48103437a98322c491 Mon Sep 17 00:00:00 2001
From: clubby789 <jamie@hill-daniel.co.uk>
Date: Thu, 3 Oct 2024 15:10:47 +0000
Subject: [PATCH] JumpThreading: fix bitwise Not on non-booleans

---
 compiler/rustc_middle/src/ty/consts/int.rs    |  5 +++++
 .../rustc_mir_transform/src/jump_threading.rs |  6 ++----
 ...ing.not_int.JumpThreading.panic-abort.diff |  2 +-
 ...ng.not_int.JumpThreading.panic-unwind.diff |  2 +-
 tests/mir-opt/jump_threading.rs               |  6 ++++++
 ..._forward.PreCodegen.after.panic-unwind.mir | 19 +++++++++++--------
 6 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/compiler/rustc_middle/src/ty/consts/int.rs b/compiler/rustc_middle/src/ty/consts/int.rs
index 1732374ab8c26..a278fc2e7fa4f 100644
--- a/compiler/rustc_middle/src/ty/consts/int.rs
+++ b/compiler/rustc_middle/src/ty/consts/int.rs
@@ -398,6 +398,11 @@ impl ScalarInt {
     pub fn to_f128(self) -> Quad {
         self.to_float()
     }
+
+    #[inline]
+    pub fn bitwise_not(self) -> Self {
+        Self { data: self.size().truncate(!self.data), size: self.size }
+    }
 }
 
 macro_rules! from_x_for_scalar_int {
diff --git a/compiler/rustc_mir_transform/src/jump_threading.rs b/compiler/rustc_mir_transform/src/jump_threading.rs
index 1844b97887ab9..0c0a354b48128 100644
--- a/compiler/rustc_mir_transform/src/jump_threading.rs
+++ b/compiler/rustc_mir_transform/src/jump_threading.rs
@@ -150,10 +150,8 @@ impl Condition {
     }
 
     fn inv(mut self) -> Self {
-        self.polarity = match self.polarity {
-            Polarity::Eq => Polarity::Ne,
-            Polarity::Ne => Polarity::Eq,
-        };
+        self.value = self.value.bitwise_not();
+
         self
     }
 }
diff --git a/tests/mir-opt/jump_threading.not_int.JumpThreading.panic-abort.diff b/tests/mir-opt/jump_threading.not_int.JumpThreading.panic-abort.diff
index e2d1e295f533a..beb5cd59c8a15 100644
--- a/tests/mir-opt/jump_threading.not_int.JumpThreading.panic-abort.diff
+++ b/tests/mir-opt/jump_threading.not_int.JumpThreading.panic-abort.diff
@@ -22,7 +22,7 @@
           StorageDead(_4);
           _2 = Eq(move _3, const 0_i32);
 -         switchInt(move _2) -> [0: bb2, otherwise: bb1];
-+         goto -> bb1;
++         goto -> bb2;
       }
   
       bb1: {
diff --git a/tests/mir-opt/jump_threading.not_int.JumpThreading.panic-unwind.diff b/tests/mir-opt/jump_threading.not_int.JumpThreading.panic-unwind.diff
index e2d1e295f533a..beb5cd59c8a15 100644
--- a/tests/mir-opt/jump_threading.not_int.JumpThreading.panic-unwind.diff
+++ b/tests/mir-opt/jump_threading.not_int.JumpThreading.panic-unwind.diff
@@ -22,7 +22,7 @@
           StorageDead(_4);
           _2 = Eq(move _3, const 0_i32);
 -         switchInt(move _2) -> [0: bb2, otherwise: bb1];
-+         goto -> bb1;
++         goto -> bb2;
       }
   
       bb1: {
diff --git a/tests/mir-opt/jump_threading.rs b/tests/mir-opt/jump_threading.rs
index f62461e225b43..1a1f658640581 100644
--- a/tests/mir-opt/jump_threading.rs
+++ b/tests/mir-opt/jump_threading.rs
@@ -533,7 +533,13 @@ fn floats() -> u32 {
 
 fn not_int() -> i32 {
     // CHECK-LABEL: fn not_int(
+    // CHECK: goto -> bb2
 
+    // CHECK-LABEL: bb1: {
+    // _0 = const 1_i32;
+
+    // CHECK-LABEL: bb2: {
+    // _0 = const 0_i32;
     // Test for issue #131195, where !a == b is assumed to be equivalent to a != b
     // This is only the case for bools
     let a = 1;
diff --git a/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir b/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir
index bf1ffd1ef328d..9fbfd99a3850f 100644
--- a/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir
+++ b/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir
@@ -33,7 +33,7 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
         StorageLive(_4);
         StorageLive(_3);
         _3 = Gt(copy _2, const 65535_usize);
-        switchInt(move _3) -> [0: bb1, otherwise: bb5];
+        switchInt(move _3) -> [0: bb1, otherwise: bb6];
     }
 
     bb1: {
@@ -54,27 +54,30 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
     bb3: {
         StorageDead(_5);
         StorageDead(_6);
-        StorageDead(_7);
-        goto -> bb7;
+        goto -> bb5;
     }
 
     bb4: {
         StorageDead(_5);
         StorageDead(_6);
-        StorageDead(_7);
-        goto -> bb6;
+        goto -> bb5;
     }
 
     bb5: {
-        StorageDead(_3);
-        goto -> bb6;
+        StorageDead(_7);
+        goto -> bb7;
     }
 
     bb6: {
-        assert(!const true, "attempt to compute `{} + {}`, which would overflow", const core::num::<impl u16>::MAX, const 1_u16) -> [success: bb7, unwind continue];
+        StorageDead(_3);
+        goto -> bb7;
     }
 
     bb7: {
+        assert(!const true, "attempt to compute `{} + {}`, which would overflow", const core::num::<impl u16>::MAX, const 1_u16) -> [success: bb8, unwind continue];
+    }
+
+    bb8: {
         StorageLive(_8);
         _8 = copy _2 as u16 (IntToInt);
         _0 = Add(copy _1, copy _8);