Add OperandValue::Uninit to improve lowering of MaybeUninit::uninit

compiler/rustc_codegen_ssa/src/base.rs

@@ -298,7 +298,7 @@ pub(crate) fn coerce_unsized_into<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
             let (base, info) = match bx.load_operand(src).val {
                 OperandValue::Pair(base, info) => unsize_ptr(bx, base, src_ty, dst_ty, Some(info)),
                 OperandValue::Immediate(base) => unsize_ptr(bx, base, src_ty, dst_ty, None),
-                OperandValue::Ref(..) | OperandValue::ZeroSized => bug!(),
+                OperandValue::Ref(..) | OperandValue::ZeroSized | OperandValue::Uninit => bug!(),
             };
             OperandValue::Pair(base, info).store(bx, dst);
         }

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -17,7 +17,7 @@ use rustc_target::callconv::{ArgAbi, CastTarget, FnAbi, PassMode};
 use tracing::{debug, info};
 
 use super::operand::OperandRef;
-use super::operand::OperandValue::{Immediate, Pair, Ref, ZeroSized};
+use super::operand::OperandValue::{Immediate, Pair, Ref, Uninit, ZeroSized};
 use super::place::{PlaceRef, PlaceValue};
 use super::{CachedLlbb, FunctionCx, LocalRef};
 use crate::base::{self, is_call_from_compiler_builtins_to_upstream_monomorphization};
@@ -527,10 +527,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
             PassMode::Direct(_) | PassMode::Pair(..) => {
                 let op = self.codegen_consume(bx, mir::Place::return_place().as_ref());
-                if let Ref(place_val) = op.val {
-                    bx.load_from_place(bx.backend_type(op.layout), place_val)
-                } else {
-                    op.immediate_or_packed_pair(bx)
+                match op.val {
+                    Uninit => bx.const_undef(bx.immediate_backend_type(op.layout)),
+                    Ref(place_val) => bx.load_from_place(bx.backend_type(op.layout), place_val),
+                    _ => op.immediate_or_packed_pair(bx),
                 }
             }
 
@@ -557,6 +557,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         place_val.llval
                     }
                     ZeroSized => bug!("ZST return value shouldn't be in PassMode::Cast"),
+                    Uninit => {
+                        bx.ret_void();
+                        return;
+                    }
                 };
                 load_cast(bx, cast_ty, llslot, self.fn_abi.ret.layout.align.abi)
             }
@@ -1587,6 +1591,10 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 }
                 _ => bug!("ZST {op:?} wasn't ignored, but was passed with abi {arg:?}"),
             },
+            Uninit => {
+                let scratch = PlaceRef::alloca(bx, arg.layout);
+                (scratch.val.llval, scratch.val.align, true)
+            }
         };
 
         if by_ref && !arg.is_indirect() {

compiler/rustc_codegen_ssa/src/mir/debuginfo.rs

@@ -341,6 +341,9 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     OperandValue::ZeroSized => {
                         // These never have a value to talk about
                     }
+                    OperandValue::Uninit => {
+                        // Better not have a useful name
+                    }
                 },
                 LocalRef::PendingOperand => {}
             }

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -67,9 +67,14 @@ pub enum OperandValue<V> {
     /// `is_zst` on its `Layout` returns `true`. Note however that
     /// these values can still require alignment.
     ZeroSized,
+    Uninit,
 }
 
 impl<V: CodegenObject> OperandValue<V> {
+    pub(crate) fn is_uninit(&self) -> bool {
+        matches!(self, OperandValue::Uninit)
+    }
+
     /// Treat this value as a pointer and return the data pointer and
     /// optional metadata as backend values.
     ///
@@ -100,6 +105,7 @@ impl<V: CodegenObject> OperandValue<V> {
         ty: TyAndLayout<'tcx>,
     ) -> bool {
         match self {
+            OperandValue::Uninit => true,
             OperandValue::ZeroSized => ty.is_zst(),
             OperandValue::Immediate(_) => cx.is_backend_immediate(ty),
             OperandValue::Pair(_, _) => cx.is_backend_scalar_pair(ty),
@@ -144,6 +150,10 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
     ) -> Self {
         let layout = bx.layout_of(ty);
 
+        if val.all_bytes_uninit(bx.tcx()) {
+            return OperandRef { val: OperandValue::Uninit, layout };
+        }
+
         let val = match val {
             ConstValue::Scalar(x) => {
                 let BackendRepr::Scalar(scalar) = layout.backend_repr else {
@@ -442,6 +452,7 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
 
         // Read the tag/niche-encoded discriminant from memory.
         let tag_op = match self.val {
+            OperandValue::Uninit => bug!("shouldn't load from uninit"),
             OperandValue::ZeroSized => bug!(),
             OperandValue::Immediate(_) | OperandValue::Pair(_, _) => {
                 self.extract_field(fx, bx, tag_field.as_usize())
@@ -591,6 +602,28 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
 }
 
 impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, Result<V, abi::Scalar>> {
+    fn update_uninit<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
+        bx: &mut Bx,
+        tgt: &mut Result<V, abi::Scalar>,
+    ) {
+        let to_scalar = tgt.unwrap_err();
+        let bty = bx.cx().type_from_scalar(to_scalar);
+        *tgt = Ok(bx.const_undef(bty));
+    }
+
+    fn update<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
+        bx: &mut Bx,
+        tgt: &mut Result<V, abi::Scalar>,
+        src: V,
+        from_scalar: rustc_abi::Scalar,
+    ) {
+        let from_bty = bx.cx().type_from_scalar(from_scalar);
+        let to_scalar = tgt.unwrap_err();
+        let to_bty = bx.cx().type_from_scalar(to_scalar);
+        let imm = transmute_immediate(bx, src, from_scalar, from_bty, to_scalar, to_bty);
+        *tgt = Ok(imm);
+    }
+
     pub(crate) fn insert_field<Bx: BuilderMethods<'a, 'tcx, Value = V>>(
         &mut self,
         bx: &mut Bx,
@@ -614,37 +647,48 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, Result<V, abi::Scalar>> {
             (field_layout.is_zst(), field_offset == Size::ZERO)
         };
 
-        let mut update = |tgt: &mut Result<V, abi::Scalar>, src, from_scalar| {
-            let from_bty = bx.cx().type_from_scalar(from_scalar);
-            let to_scalar = tgt.unwrap_err();
-            let to_bty = bx.cx().type_from_scalar(to_scalar);
-            let imm = transmute_immediate(bx, src, from_scalar, from_bty, to_scalar, to_bty);
-            *tgt = Ok(imm);
-        };
-
         match (operand.val, operand.layout.backend_repr) {
             (OperandValue::ZeroSized, _) if expect_zst => {}
             (OperandValue::Immediate(v), BackendRepr::Scalar(from_scalar)) => match &mut self.val {
                 OperandValue::Immediate(val @ Err(_)) if is_zero_offset => {
-                    update(val, v, from_scalar);
+                    Self::update(bx, val, v, from_scalar);
                 }
                 OperandValue::Pair(fst @ Err(_), _) if is_zero_offset => {
-                    update(fst, v, from_scalar);
+                    Self::update(bx, fst, v, from_scalar);
                 }
                 OperandValue::Pair(_, snd @ Err(_)) if !is_zero_offset => {
-                    update(snd, v, from_scalar);
+                    Self::update(bx, snd, v, from_scalar);
+                }
+                _ => bug!("Tried to insert {operand:?} into {v:?}.{f:?} of {self:?}"),
+            },
+            (OperandValue::Uninit, BackendRepr::Scalar(_)) => match &mut self.val {
+                OperandValue::Immediate(val @ Err(_)) if is_zero_offset => {
+                    Self::update_uninit(bx, val);
+                }
+                OperandValue::Pair(fst @ Err(_), _) if is_zero_offset => {
+                    Self::update_uninit(bx, fst);
+                }
+                OperandValue::Pair(_, snd @ Err(_)) if !is_zero_offset => {
+                    Self::update_uninit(bx, snd);
                 }
                 _ => bug!("Tried to insert {operand:?} into {v:?}.{f:?} of {self:?}"),
             },
             (OperandValue::Pair(a, b), BackendRepr::ScalarPair(from_sa, from_sb)) => {
                 match &mut self.val {
                     OperandValue::Pair(fst @ Err(_), snd @ Err(_)) => {
-                        update(fst, a, from_sa);
-                        update(snd, b, from_sb);
+                        Self::update(bx, fst, a, from_sa);
+                        Self::update(bx, snd, b, from_sb);
                     }
                     _ => bug!("Tried to insert {operand:?} into {v:?}.{f:?} of {self:?}"),
                 }
             }
+            (OperandValue::Uninit, BackendRepr::ScalarPair(..)) => match &mut self.val {
+                OperandValue::Pair(fst @ Err(_), snd @ Err(_)) => {
+                    Self::update_uninit(bx, fst);
+                    Self::update_uninit(bx, snd);
+                }
+                _ => bug!("Tried to insert {operand:?} into {v:?}.{f:?} of {self:?}"),
+            },
             _ => bug!("Unsupported operand {operand:?} inserting into {v:?}.{f:?} of {self:?}"),
         }
     }
@@ -663,6 +707,7 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, Result<V, abi::Scalar>> {
         };
 
         let val = match val {
+            OperandValue::Uninit => OperandValue::Uninit,
             OperandValue::ZeroSized => OperandValue::ZeroSized,
             OperandValue::Immediate(v) => OperandValue::Immediate(unwrap(v)),
             OperandValue::Pair(a, b) => OperandValue::Pair(unwrap(a), unwrap(b)),
@@ -739,6 +784,13 @@ impl<'a, 'tcx, V: CodegenObject> OperandValue<V> {
     ) {
         debug!("OperandRef::store: operand={:?}, dest={:?}", self, dest);
         match self {
+            OperandValue::Uninit => {
+                // Ideally we'd hint to the backend that the destination is deinitialized by the
+                // store. But in practice the destination is almost always uninit already because
+                // OperandValue::Uninit is pretty much only produced by MaybeUninit::uninit.
+                // Attempting to generate a hint by calling memset with undef mostly seems to
+                // confuse LLVM.
+            }
             OperandValue::ZeroSized => {
                 // Avoid generating stores of zero-sized values, because the only way to have a
                 // zero-sized value is through `undef`/`poison`, and the store itself is useless.

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -67,6 +67,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                         base::coerce_unsized_into(bx, scratch, dest);
                         scratch.storage_dead(bx);
                     }
+                    OperandValue::Uninit => {}
                     OperandValue::Ref(val) => {
                         if val.llextra.is_some() {
                             bug!("unsized coercion on an unsized rvalue");
@@ -90,24 +91,13 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     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);
+                // Normally the check for uninit is handled inside the operand helpers, but in this
+                // one case we want to bail early so that we don't generate the loop form with an
+                // empty body.
+                if cg_elem.val.is_uninit() {
+                    return;
+                }
 
                 let try_init_all_same = |bx: &mut Bx, v| {
                     let start = dest.val.llval;
@@ -206,7 +196,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         }
 
         match src.val {
-            OperandValue::Ref(..) | OperandValue::ZeroSized => {
+            OperandValue::Ref(..) | OperandValue::ZeroSized | OperandValue::Uninit => {
                 span_bug!(
                     self.mir.span,
                     "Operand path should have handled transmute \
@@ -251,6 +241,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         let cast_kind = self.value_kind(cast);
 
         match operand.val {
+            OperandValue::Uninit => Some(OperandValue::Uninit),
             OperandValue::Ref(source_place_val) => {
                 assert_eq!(source_place_val.llextra, None);
                 assert_matches!(operand_kind, OperandValueKind::Ref);

library/core/src/mem/maybe_uninit.rs

@@ -328,7 +328,10 @@ impl<T> MaybeUninit<T> {
     #[inline(always)]
     #[rustc_diagnostic_item = "maybe_uninit_uninit"]
     pub const fn uninit() -> MaybeUninit<T> {
-        MaybeUninit { uninit: () }
+        // It is very helpful for codegen to know when are writing uninit bytes. MIR optimizations
+        // currently do not const-propagate unions, but if we create the const manually that can be
+        // trivially propagated. See #142837.
+        const { MaybeUninit { uninit: () } }
     }
 
     /// Creates a new `MaybeUninit<T>` in an uninitialized state, with the memory being

tests/codegen/maybeuninit.rs

@@ -0,0 +1,32 @@
+//@ compile-flags: -Copt-level=3 -Cdebuginfo=0
+
+// This is a regression test for https://github.com/rust-lang/rust/issues/139355 as well as
+// regressions I introduced while implementing a solution.
+
+#![crate_type = "lib"]
+
+use std::mem::MaybeUninit;
+
+// CHECK-LABEL: @create_small_uninit_array
+#[no_mangle]
+fn create_small_uninit_array() -> [MaybeUninit<u8>; 4] {
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: ret i32 undef
+    [MaybeUninit::<u8>::uninit(); 4]
+}
+
+// CHECK-LABEL: @create_nested_uninit_array
+#[no_mangle]
+fn create_nested_uninit_array() -> [[MaybeUninit<u8>; 4]; 100] {
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: ret void
+    [[MaybeUninit::<u8>::uninit(); 4]; 100]
+}
+
+// CHECK-LABEL: @create_ptr
+#[no_mangle]
+fn create_ptr() -> MaybeUninit<&'static str> {
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: ret { ptr, i64 } undef
+    MaybeUninit::uninit()
+}

tests/codegen/uninit-consts.rs

@@ -11,21 +11,18 @@ pub struct PartiallyUninit {
     y: MaybeUninit<[u8; 10]>,
 }
 
-// CHECK: [[FULLY_UNINIT:@.*]] = private unnamed_addr constant [10 x i8] undef
-
 // CHECK: [[PARTIALLY_UNINIT:@.*]] = private unnamed_addr constant <{ [4 x i8], [12 x i8] }> <{ [4 x i8] c"{{\\EF\\BE\\AD\\DE|\\DE\\AD\\BE\\EF}}", [12 x i8] undef }>, align 4
 
 // This shouldn't contain undef, since it contains more chunks
 // than the default value of uninit_const_chunk_threshold.
 // CHECK: [[UNINIT_PADDING_HUGE:@.*]] = private unnamed_addr constant [32768 x i8] c"{{.+}}", align 4
 
-// CHECK: [[FULLY_UNINIT_HUGE:@.*]] = private unnamed_addr constant [16384 x i8] undef
-
 // CHECK-LABEL: @fully_uninit
 #[no_mangle]
 pub const fn fully_uninit() -> MaybeUninit<[u8; 10]> {
     const M: MaybeUninit<[u8; 10]> = MaybeUninit::uninit();
-    // CHECK: call void @llvm.memcpy.{{.+}}(ptr align 1 %_0, ptr align 1 {{.*}}[[FULLY_UNINIT]]{{.*}}, i{{(32|64)}} 10, i1 false)
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: ret void
     M
 }
 
@@ -49,6 +46,7 @@ pub const fn uninit_padding_huge() -> [(u32, u8); 4096] {
 #[no_mangle]
 pub const fn fully_uninit_huge() -> MaybeUninit<[u32; 4096]> {
     const F: MaybeUninit<[u32; 4096]> = MaybeUninit::uninit();
-    // CHECK: call void @llvm.memcpy.{{.+}}(ptr align 4 %_0, ptr align 4 {{.*}}[[FULLY_UNINIT_HUGE]]{{.*}}, i{{(32|64)}} 16384, i1 false)
+    // CHECK-NEXT: start:
+    // CHECK-NEXT: ret void
     F
 }