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

7 files changed, 102 insertions, 23 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..e37ead3674b 100644
--- a/compiler/rustc_mir_transform/src/remove_zsts.rs
+++ b/compiler/rustc_mir_transform/src/remove_zsts.rs
@@ -36,31 +36,37 @@ 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() {
+        // definitely ZST
+        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),
         // 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,
+        | ty::Alias(ty::Opaque, ..) => None,
         // unreachable or can't be ZST
-        _ => false,
+        _ => Some(false),
     }
 }
 
 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..a5e3ddbc57e
--- /dev/null
+++ b/tests/mir-opt/instsimplify/simplify_repeat.repeat_once_to_aggregate.InstSimplify-after-simplifycfg.diff
@@ -0,0 +1,18 @@
+- // 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 mut _2: T;
+  
+      bb0: {
+          StorageLive(_2);
+          _2 = copy _1;
+-         _0 = [move _2; 1];
++         _0 = [move _2];
+          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..359f66710b2
--- /dev/null
+++ b/tests/mir-opt/instsimplify/simplify_repeat.rs
@@ -0,0 +1,13 @@
+//@ test-mir-pass: InstSimplify-after-simplifycfg
+//@ compile-flags: -C panic=abort
+#![crate_type = "lib"]
+
+// 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-NOT: [move {{_[0-9]+}}; 1]
+    // CHECK: _0 = [move {{_[0-9]+}}];
+    // CHECK-NOT: [move {{_[0-9]+}}; 1]
+
+    [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..97ff88cc812
--- /dev/null
+++ b/tests/mir-opt/remove_zsts.remove_generic_array.RemoveZsts.diff
@@ -0,0 +1,29 @@
+- // 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;
+      scope 1 {
+-         debug a => _2;
++         debug a => const ZeroSized: [T; 0];
+      }
+  
+      bb0: {
+-         StorageLive(_2);
++         nop;
+          StorageLive(_3);
+          _3 = copy _1;
+-         _2 = [];
++         nop;
+          StorageDead(_3);
+-         _0 = const ();
+-         StorageDead(_2);
++         nop;
++         nop;
+          return;
+      }
+  }
+  
diff --git a/tests/mir-opt/remove_zsts.rs b/tests/mir-opt/remove_zsts.rs
index e33a272fe16..d3db1c20142 100644
--- a/tests/mir-opt/remove_zsts.rs
+++ b/tests/mir-opt/remove_zsts.rs
@@ -1,15 +1,25 @@
-// 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: () }
 }
 
+// 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-NOT: = [];
+    let a = [x; 0];
+}
+
 fn main() {
     get_union();
 }