Auto merge of #135322 - scottmcm:inst-simplify-repeat-one, r=oli-obk

[mir-opt] simplify `Repeat`s that don't actually repeat the operand

Created because when I was writing this case in GVN

https://github.com/rust-lang/rust/pull/133324/files#diff-292b215fdc6b59e3f3c4173c2270b14591c5673832fbfb05cd69a05c2ef0c30eR977-R979

I happened to notice that it worked for `[x]` but not for `[x; 1]`, so figured it would be good to simplify that `Repeat` to the simpler `Aggregate`.

7 files changed, 149 insertions, 29 deletions

diff --git a/compiler/rustc_mir_transform/src/instsimplify.rs b/compiler/rustc_mir_transform/src/instsimplify.rs
index f74a577c598..20e2e3e8ba2 100644
--- a/compiler/rustc_mir_transform/src/instsimplify.rs
+++ b/compiler/rustc_mir_transform/src/instsimplify.rs
@@ -49,6 +49,7 @@ impl<'tcx> crate::MirPass<'tcx> for InstSimplify {
                         ctx.simplify_ptr_aggregate(rvalue);
                         ctx.simplify_cast(rvalue);
                         ctx.simplify_repeated_aggregate(rvalue);
+                        ctx.simplify_repeat_once(rvalue);
                     }
                     _ => {}
                 }
@@ -207,6 +208,18 @@ impl<'tcx> InstSimplifyContext<'_, 'tcx> {
         }
     }
 
+    /// Simplify `[x; 1]` to just `[x]`.
+    fn simplify_repeat_once(&self, rvalue: &mut Rvalue<'tcx>) {
+        if let Rvalue::Repeat(operand, count) = rvalue
+            && let Some(1) = count.try_to_target_usize(self.tcx)
+        {
+            *rvalue = Rvalue::Aggregate(
+                Box::new(AggregateKind::Array(operand.ty(self.local_decls, self.tcx))),
+                [operand.clone()].into(),
+            );
+        }
+    }
+
     fn simplify_primitive_clone(
         &self,
         terminator: &mut Terminator<'tcx>,
diff --git a/compiler/rustc_mir_transform/src/remove_zsts.rs b/compiler/rustc_mir_transform/src/remove_zsts.rs
index 64e183bcbc0..55e5701bd0a 100644
--- a/compiler/rustc_mir_transform/src/remove_zsts.rs
+++ b/compiler/rustc_mir_transform/src/remove_zsts.rs
@@ -36,31 +36,39 @@ struct Replacer<'a, 'tcx> {
 }
 
 /// A cheap, approximate check to avoid unnecessary `layout_of` calls.
-fn maybe_zst(ty: Ty<'_>) -> bool {
+///
+/// `Some(true)` is definitely ZST; `Some(false)` is definitely *not* ZST.
+///
+/// `None` may or may not be, and must check `layout_of` to be sure.
+fn trivially_zst<'tcx>(ty: Ty<'tcx>, tcx: TyCtxt<'tcx>) -> Option<bool> {
     match ty.kind() {
-        // maybe ZST (could be more precise)
-        ty::Adt(..)
-        | ty::Array(..)
-        | ty::Closure(..)
-        | ty::CoroutineClosure(..)
-        | ty::Tuple(..)
-        | ty::Alias(ty::Opaque, ..) => true,
         // definitely ZST
-        ty::FnDef(..) | ty::Never => true,
-        // unreachable or can't be ZST
-        _ => false,
+        ty::FnDef(..) | ty::Never => Some(true),
+        ty::Tuple(fields) if fields.is_empty() => Some(true),
+        ty::Array(_ty, len) if let Some(0) = len.try_to_target_usize(tcx) => Some(true),
+        // clearly not ZST
+        ty::Bool
+        | ty::Char
+        | ty::Int(..)
+        | ty::Uint(..)
+        | ty::Float(..)
+        | ty::RawPtr(..)
+        | ty::Ref(..)
+        | ty::FnPtr(..) => Some(false),
+        // check `layout_of` to see (including unreachable things we won't actually see)
+        _ => None,
     }
 }
 
 impl<'tcx> Replacer<'_, 'tcx> {
     fn known_to_be_zst(&self, ty: Ty<'tcx>) -> bool {
-        if !maybe_zst(ty) {
-            return false;
+        if let Some(is_zst) = trivially_zst(ty, self.tcx) {
+            is_zst
+        } else {
+            self.tcx
+                .layout_of(self.typing_env.as_query_input(ty))
+                .is_ok_and(|layout| layout.is_zst())
         }
-        let Ok(layout) = self.tcx.layout_of(self.typing_env.as_query_input(ty)) else {
-            return false;
-        };
-        layout.is_zst()
     }
 
     fn make_zst(&self, ty: Ty<'tcx>) -> ConstOperand<'tcx> {
diff --git a/tests/mir-opt/instsimplify/simplify_repeat.repeat_once_to_aggregate.InstSimplify-after-simplifycfg.diff b/tests/mir-opt/instsimplify/simplify_repeat.repeat_once_to_aggregate.InstSimplify-after-simplifycfg.diff
new file mode 100644
index 00000000000..6c1b9abc5d7
--- /dev/null
+++ b/tests/mir-opt/instsimplify/simplify_repeat.repeat_once_to_aggregate.InstSimplify-after-simplifycfg.diff
@@ -0,0 +1,30 @@
+- // MIR for `repeat_once_to_aggregate` before InstSimplify-after-simplifycfg
++ // MIR for `repeat_once_to_aggregate` after InstSimplify-after-simplifycfg
+  
+  fn repeat_once_to_aggregate(_1: T) -> [T; 1] {
+      debug x => _1;
+      let mut _0: [T; 1];
+      let _2: [T; 1];
+      let mut _3: T;
+      let mut _4: T;
+      scope 1 {
+          debug other => _2;
+      }
+  
+      bb0: {
+          StorageLive(_2);
+          StorageLive(_3);
+          _3 = copy _1;
+-         _2 = [move _3; 1];
++         _2 = [move _3];
+          StorageDead(_3);
+          StorageLive(_4);
+          _4 = copy _1;
+-         _0 = [move _4; 1];
++         _0 = [move _4];
+          StorageDead(_4);
+          StorageDead(_2);
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/instsimplify/simplify_repeat.rs b/tests/mir-opt/instsimplify/simplify_repeat.rs
new file mode 100644
index 00000000000..abcdf32072b
--- /dev/null
+++ b/tests/mir-opt/instsimplify/simplify_repeat.rs
@@ -0,0 +1,20 @@
+//@ test-mir-pass: InstSimplify-after-simplifycfg
+//@ compile-flags: -C panic=abort
+#![crate_type = "lib"]
+
+const MYSTERY: usize = 3_usize.pow(2) - 2_usize.pow(3);
+
+// EMIT_MIR simplify_repeat.repeat_once_to_aggregate.InstSimplify-after-simplifycfg.diff
+pub fn repeat_once_to_aggregate<T: Copy>(x: T) -> [T; 1] {
+    // CHECK-LABEL: fn repeat_once_to_aggregate(
+    // CHECK: debug other => [[OTHER:_[0-9]+]]
+    // CHECK-NOT: [move {{_[0-9]+}}; 1]
+    // CHECK: [[OTHER]] = [move {{_[0-9]+}}];
+    // CHECK-NOT: [move {{_[0-9]+}}; 1]
+    // CHECK: _0 = [move {{_[0-9]+}}];
+    // CHECK-NOT: [move {{_[0-9]+}}; 1]
+
+    let other = [x; MYSTERY];
+
+    [x; 1]
+}
diff --git a/tests/mir-opt/remove_zsts.get_union.PreCodegen.after.mir b/tests/mir-opt/remove_zsts.get_union.PreCodegen.after.mir
deleted file mode 100644
index 5886a5bfeea..00000000000
--- a/tests/mir-opt/remove_zsts.get_union.PreCodegen.after.mir
+++ /dev/null
@@ -1,10 +0,0 @@
-// MIR for `get_union` after PreCodegen
-
-fn get_union() -> Foo {
-    let mut _0: Foo;
-
-    bb0: {
-        _0 = Foo { x: const () };
-        return;
-    }
-}
diff --git a/tests/mir-opt/remove_zsts.remove_generic_array.RemoveZsts.diff b/tests/mir-opt/remove_zsts.remove_generic_array.RemoveZsts.diff
new file mode 100644
index 00000000000..2ac944a6c6b
--- /dev/null
+++ b/tests/mir-opt/remove_zsts.remove_generic_array.RemoveZsts.diff
@@ -0,0 +1,44 @@
+- // MIR for `remove_generic_array` before RemoveZsts
++ // MIR for `remove_generic_array` after RemoveZsts
+  
+  fn remove_generic_array(_1: T) -> () {
+      debug x => _1;
+      let mut _0: ();
+      let _2: [T; 0];
+      let mut _3: T;
+      let mut _5: T;
+      scope 1 {
+-         debug a => _2;
++         debug a => const ZeroSized: [T; 0];
+          let _4: [T; 0];
+          scope 2 {
+-             debug b => _4;
++             debug b => const ZeroSized: [T; 0];
+          }
+      }
+  
+      bb0: {
+-         StorageLive(_2);
++         nop;
+          StorageLive(_3);
+          _3 = copy _1;
+-         _2 = [];
++         nop;
+          StorageDead(_3);
+-         StorageLive(_4);
++         nop;
+          StorageLive(_5);
+          _5 = copy _1;
+-         _4 = [];
++         nop;
+          StorageDead(_5);
+-         _0 = const ();
+-         StorageDead(_4);
+-         StorageDead(_2);
++         nop;
++         nop;
++         nop;
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/remove_zsts.rs b/tests/mir-opt/remove_zsts.rs
index e33a272fe16..baf9d8ece26 100644
--- a/tests/mir-opt/remove_zsts.rs
+++ b/tests/mir-opt/remove_zsts.rs
@@ -1,15 +1,30 @@
-// skip-filecheck
+//@ test-mir-pass: RemoveZsts
+
 union Foo {
     x: (),
     y: u64,
 }
 
 // EMIT_MIR remove_zsts.get_union.RemoveZsts.diff
-// EMIT_MIR remove_zsts.get_union.PreCodegen.after.mir
 fn get_union() -> Foo {
+    // CHECK-LABEL: fn get_union
+    // CHECK: _0 = Foo { x: const () };
     Foo { x: () }
 }
 
+const MYSTERY: usize = 280_usize.isqrt() - 260_usize.isqrt();
+
+// EMIT_MIR remove_zsts.remove_generic_array.RemoveZsts.diff
+fn remove_generic_array<T: Copy>(x: T) {
+    // CHECK-LABEL: fn remove_generic_array
+    // CHECK: debug a => const ZeroSized: [T; 0];
+    // CHECK: debug b => const ZeroSized: [T; 0];
+    // CHECK-NOT: = [];
+    // CHECK-NOT: ; 1]
+    let a = [x; 0];
+    let b = [x; MYSTERY];
+}
+
 fn main() {
     get_union();
 }