From 9d266bf34fcea419512a34505dc4cc565aa47ba6 Mon Sep 17 00:00:00 2001
From: Scott McMurray <scottmcm@users.noreply.github.com>
Date: Sun, 12 May 2024 00:22:42 -0700
Subject: [PATCH] Enable GVN for `AggregateKind::RawPtr`

---
 compiler/rustc_mir_transform/src/gvn.rs       | 11 ++++-
 tests/mir-opt/gvn.rs                          | 12 +++++
 .../gvn.slice_index.GVN.panic-abort.diff      | 48 +++++++++++++++++++
 .../gvn.slice_index.GVN.panic-unwind.diff     | 48 +++++++++++++++++++
 4 files changed, 117 insertions(+), 2 deletions(-)
 create mode 100644 tests/mir-opt/gvn.slice_index.GVN.panic-abort.diff
 create mode 100644 tests/mir-opt/gvn.slice_index.GVN.panic-unwind.diff

diff --git a/compiler/rustc_mir_transform/src/gvn.rs b/compiler/rustc_mir_transform/src/gvn.rs
index 342d1a1cd5bb0..b10035688ce6e 100644
--- a/compiler/rustc_mir_transform/src/gvn.rs
+++ b/compiler/rustc_mir_transform/src/gvn.rs
@@ -176,6 +176,9 @@ enum AggregateTy<'tcx> {
     Array,
     Tuple,
     Def(DefId, ty::GenericArgsRef<'tcx>),
+    /// The type of the pointer can't be determined from the operands,
+    /// so we have to keep the full thing here.
+    RawPtr(Ty<'tcx>),
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
@@ -385,6 +388,7 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
                     AggregateTy::Def(def_id, args) => {
                         self.tcx.type_of(def_id).instantiate(self.tcx, args)
                     }
+                    AggregateTy::RawPtr(ty) => ty,
                 };
                 let variant = if ty.is_enum() { Some(variant) } else { None };
                 let ty = self.ecx.layout_of(ty).ok()?;
@@ -927,8 +931,11 @@ impl<'body, 'tcx> VnState<'body, 'tcx> {
             }
             // Do not track unions.
             AggregateKind::Adt(_, _, _, _, Some(_)) => return None,
-            // FIXME: Do the extra work to GVN `from_raw_parts`
-            AggregateKind::RawPtr(..) => return None,
+            AggregateKind::RawPtr(pty, mtbl) => {
+                assert_eq!(fields.len(), 2);
+                let ty = Ty::new_ptr(self.tcx, pty, mtbl);
+                (AggregateTy::RawPtr(ty), FIRST_VARIANT)
+            }
         };
 
         let fields: Option<Vec<_>> = fields
diff --git a/tests/mir-opt/gvn.rs b/tests/mir-opt/gvn.rs
index 9be3051528356..69e4b84bb1f29 100644
--- a/tests/mir-opt/gvn.rs
+++ b/tests/mir-opt/gvn.rs
@@ -781,6 +781,16 @@ fn non_freeze<T: Copy>(x: T) {
     )
 }
 
+// Check that we can const-prop into `from_raw_parts`
+fn slice_index(x: &[i32]) -> *const [i32] {
+    // CHECK-LABEL: fn slice_index(
+    // CHECK: _0 = *const [i32] from
+    // CHECK-SAME: const 123_usize);
+    let ptr = x.as_ptr();
+    let len = 123;
+    std::intrinsics::aggregate_raw_ptr(ptr, len)
+}
+
 fn main() {
     subexpression_elimination(2, 4, 5);
     wrap_unwrap(5);
@@ -805,6 +815,7 @@ fn main() {
     wide_ptr_integer();
     borrowed(5);
     non_freeze(5);
+    slice_index(&[1]);
 }
 
 #[inline(never)]
@@ -838,3 +849,4 @@ fn identity<T>(x: T) -> T {
 // EMIT_MIR gvn.wide_ptr_integer.GVN.diff
 // EMIT_MIR gvn.borrowed.GVN.diff
 // EMIT_MIR gvn.non_freeze.GVN.diff
+// EMIT_MIR gvn.slice_index.GVN.diff
diff --git a/tests/mir-opt/gvn.slice_index.GVN.panic-abort.diff b/tests/mir-opt/gvn.slice_index.GVN.panic-abort.diff
new file mode 100644
index 0000000000000..4349b8e937f33
--- /dev/null
+++ b/tests/mir-opt/gvn.slice_index.GVN.panic-abort.diff
@@ -0,0 +1,48 @@
+- // MIR for `slice_index` before GVN
++ // MIR for `slice_index` after GVN
+  
+  fn slice_index(_1: &[i32]) -> *const [i32] {
+      debug x => _1;
+      let mut _0: *const [i32];
+      let _2: *const i32;
+      let mut _3: &[i32];
+      let mut _5: *const i32;
+      let mut _6: usize;
+      scope 1 {
+          debug ptr => _2;
+          let _4: usize;
+          scope 2 {
+              debug len => _4;
+          }
+      }
+  
+      bb0: {
+-         StorageLive(_2);
++         nop;
+          StorageLive(_3);
+          _3 = &(*_1);
+          _2 = core::slice::<impl [i32]>::as_ptr(move _3) -> [return: bb1, unwind unreachable];
+      }
+  
+      bb1: {
+          StorageDead(_3);
+-         StorageLive(_4);
++         nop;
+          _4 = const 123_usize;
+          StorageLive(_5);
+          _5 = _2;
+          StorageLive(_6);
+-         _6 = _4;
+-         _0 = *const [i32] from (move _5, move _6);
++         _6 = const 123_usize;
++         _0 = *const [i32] from (_2, const 123_usize);
+          StorageDead(_6);
+          StorageDead(_5);
+-         StorageDead(_4);
+-         StorageDead(_2);
++         nop;
++         nop;
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/gvn.slice_index.GVN.panic-unwind.diff b/tests/mir-opt/gvn.slice_index.GVN.panic-unwind.diff
new file mode 100644
index 0000000000000..1aa67f2064b28
--- /dev/null
+++ b/tests/mir-opt/gvn.slice_index.GVN.panic-unwind.diff
@@ -0,0 +1,48 @@
+- // MIR for `slice_index` before GVN
++ // MIR for `slice_index` after GVN
+  
+  fn slice_index(_1: &[i32]) -> *const [i32] {
+      debug x => _1;
+      let mut _0: *const [i32];
+      let _2: *const i32;
+      let mut _3: &[i32];
+      let mut _5: *const i32;
+      let mut _6: usize;
+      scope 1 {
+          debug ptr => _2;
+          let _4: usize;
+          scope 2 {
+              debug len => _4;
+          }
+      }
+  
+      bb0: {
+-         StorageLive(_2);
++         nop;
+          StorageLive(_3);
+          _3 = &(*_1);
+          _2 = core::slice::<impl [i32]>::as_ptr(move _3) -> [return: bb1, unwind continue];
+      }
+  
+      bb1: {
+          StorageDead(_3);
+-         StorageLive(_4);
++         nop;
+          _4 = const 123_usize;
+          StorageLive(_5);
+          _5 = _2;
+          StorageLive(_6);
+-         _6 = _4;
+-         _0 = *const [i32] from (move _5, move _6);
++         _6 = const 123_usize;
++         _0 = *const [i32] from (_2, const 123_usize);
+          StorageDead(_6);
+          StorageDead(_5);
+-         StorageDead(_4);
+-         StorageDead(_2);
++         nop;
++         nop;
+          return;
+      }
+  }
+