From d0986f45e094b011d132238563b05c09e4f1e20e Mon Sep 17 00:00:00 2001
From: DianQK <dianqk@dianqk.net>
Date: Mon, 9 Dec 2024 20:25:41 +0800
Subject: [PATCH] dataflow_const_prop: do not eval a ptr address in SwitchInt

---
 .../src/dataflow_const_prop.rs                |  9 +++++++--
 tests/crashes/131227.rs                       | 16 ----------------
 .../ptr-in-switch-int-issue-131227.rs         | 19 +++++++++++++++++++
 3 files changed, 26 insertions(+), 18 deletions(-)
 delete mode 100644 tests/crashes/131227.rs
 create mode 100644 tests/ui/dataflow_const_prop/ptr-in-switch-int-issue-131227.rs

diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
index d017202f48bb9..b94c925b1dbd7 100644
--- a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
+++ b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -534,8 +534,13 @@ impl<'a, 'tcx> ConstAnalysis<'a, 'tcx> {
             // This allows the set of visited edges to grow monotonically with the lattice.
             FlatSet::Bottom => TerminatorEdges::None,
             FlatSet::Elem(scalar) => {
-                let choice = scalar.assert_scalar_int().to_bits_unchecked();
-                TerminatorEdges::Single(targets.target_for_value(choice))
+                if let Ok(scalar_int) = scalar.try_to_scalar_int() {
+                    TerminatorEdges::Single(
+                        targets.target_for_value(scalar_int.to_bits_unchecked()),
+                    )
+                } else {
+                    TerminatorEdges::SwitchInt { discr, targets }
+                }
             }
             FlatSet::Top => TerminatorEdges::SwitchInt { discr, targets },
         }
diff --git a/tests/crashes/131227.rs b/tests/crashes/131227.rs
deleted file mode 100644
index f46185b5b4a66..0000000000000
--- a/tests/crashes/131227.rs
+++ /dev/null
@@ -1,16 +0,0 @@
-//@ known-bug: #131227
-//@ compile-flags: -Zmir-opt-level=3
-
-static mut G: () = ();
-
-fn myfunc() -> i32 {
-    let var = &raw mut G;
-    if var.is_null() {
-        return 0;
-    }
-    0
-}
-
-fn main() {
-    myfunc();
-}
diff --git a/tests/ui/dataflow_const_prop/ptr-in-switch-int-issue-131227.rs b/tests/ui/dataflow_const_prop/ptr-in-switch-int-issue-131227.rs
new file mode 100644
index 0000000000000..7a55e13d0eed0
--- /dev/null
+++ b/tests/ui/dataflow_const_prop/ptr-in-switch-int-issue-131227.rs
@@ -0,0 +1,19 @@
+//! Issue: <https://github.com/rust-lang/rust/issues/131227>
+//! Test that constant propagation in SwitchInt does not crash
+//! when encountering a ptr-to-int transmute.
+
+//@ check-pass
+//@ compile-flags: -Zmir-enable-passes=+InstSimplify-before-inline,+DataflowConstProp
+
+#![crate_type = "lib"]
+
+static mut G: i32 = 0;
+
+pub fn myfunc() -> i32 {
+    let var = &raw mut G;
+    let u: usize = unsafe { std::mem::transmute(var) };
+    match u {
+        0 => 0,
+        _ => 1,
+    }
+}