From 8e7d8ddffe802180b504e0ecaaa40b10b28b291e Mon Sep 17 00:00:00 2001
From: Ben Kimock <kimockb@gmail.com>
Date: Mon, 17 Mar 2025 21:29:07 -0400
Subject: [PATCH] Lower to a memset(undef) when Rvalue::Repeat repeats uninit

---
 compiler/rustc_codegen_ssa/src/mir/rvalue.rs | 21 +++++++++++++--
 compiler/rustc_middle/src/mir/consts.rs      | 28 ++++++++++++++++++--
 tests/codegen/uninit-repeat-in-aggregate.rs  | 21 +++++++++++++++
 3 files changed, 66 insertions(+), 4 deletions(-)
 create mode 100644 tests/codegen/uninit-repeat-in-aggregate.rs

diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
index 72cfd2bffb5..1df00465f74 100644
--- a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -86,13 +86,30 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             }
 
             mir::Rvalue::Repeat(ref elem, count) => {
-                let cg_elem = self.codegen_operand(bx, elem);
-
                 // Do not generate the loop for zero-sized elements or empty arrays.
                 if dest.layout.is_zst() {
                     return;
                 }
 
+                // When the element is a const with all bytes uninit, emit a single memset that
+                // writes undef to the entire destination.
+                if let mir::Operand::Constant(const_op) = elem {
+                    let val = self.eval_mir_constant(const_op);
+                    if val.all_bytes_uninit(self.cx.tcx()) {
+                        let size = bx.const_usize(dest.layout.size.bytes());
+                        bx.memset(
+                            dest.val.llval,
+                            bx.const_undef(bx.type_i8()),
+                            size,
+                            dest.val.align,
+                            MemFlags::empty(),
+                        );
+                        return;
+                    }
+                }
+
+                let cg_elem = self.codegen_operand(bx, elem);
+
                 let try_init_all_same = |bx: &mut Bx, v| {
                     let start = dest.val.llval;
                     let size = bx.const_usize(dest.layout.size.bytes());
diff --git a/compiler/rustc_middle/src/mir/consts.rs b/compiler/rustc_middle/src/mir/consts.rs
index 34b27c2e1cc..2b2ffa71628 100644
--- a/compiler/rustc_middle/src/mir/consts.rs
+++ b/compiler/rustc_middle/src/mir/consts.rs
@@ -9,7 +9,9 @@ use rustc_span::{DUMMY_SP, Span, Symbol};
 use rustc_type_ir::TypeVisitableExt;
 
 use super::interpret::ReportedErrorInfo;
-use crate::mir::interpret::{AllocId, ConstAllocation, ErrorHandled, Scalar, alloc_range};
+use crate::mir::interpret::{
+    AllocId, AllocRange, ConstAllocation, ErrorHandled, GlobalAlloc, Scalar, alloc_range,
+};
 use crate::mir::{Promoted, pretty_print_const_value};
 use crate::ty::print::{pretty_print_const, with_no_trimmed_paths};
 use crate::ty::{self, ConstKind, GenericArgsRef, ScalarInt, Ty, TyCtxt};
@@ -192,9 +194,31 @@ impl<'tcx> ConstValue<'tcx> {
                 .unwrap_memory()
                 .inner()
                 .provenance()
-                .range_empty(super::AllocRange::from(offset..offset + size), &tcx),
+                .range_empty(AllocRange::from(offset..offset + size), &tcx),
         }
     }
+
+    /// Check if a constant only contains uninitialized bytes.
+    pub fn all_bytes_uninit(&self, tcx: TyCtxt<'tcx>) -> bool {
+        let ConstValue::Indirect { alloc_id, .. } = self else {
+            return false;
+        };
+        let alloc = tcx.global_alloc(*alloc_id);
+        let GlobalAlloc::Memory(alloc) = alloc else {
+            return false;
+        };
+        let init_mask = alloc.0.init_mask();
+        let init_range = init_mask.is_range_initialized(AllocRange {
+            start: Size::ZERO,
+            size: Size::from_bytes(alloc.0.len()),
+        });
+        if let Err(range) = init_range {
+            if range.size == alloc.0.size() {
+                return true;
+            }
+        }
+        false
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////
diff --git a/tests/codegen/uninit-repeat-in-aggregate.rs b/tests/codegen/uninit-repeat-in-aggregate.rs
new file mode 100644
index 00000000000..0fa2eb7d56c
--- /dev/null
+++ b/tests/codegen/uninit-repeat-in-aggregate.rs
@@ -0,0 +1,21 @@
+//@ compile-flags: -Copt-level=3
+
+#![crate_type = "lib"]
+
+use std::mem::MaybeUninit;
+
+// We need to make sure len is at offset 0, otherwise codegen needs an extra instruction
+#[repr(C)]
+pub struct SmallVec<T> {
+    pub len: u64,
+    pub arr: [MaybeUninit<T>; 24],
+}
+
+// CHECK-LABEL: @uninit_arr_via_const
+#[no_mangle]
+pub fn uninit_arr_via_const() -> SmallVec<String> {
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: store i64 0,
+    // CHECK-NEXT: ret
+    SmallVec { len: 0, arr: [const { MaybeUninit::uninit() }; 24] }
+}