Rollup merge of #83698 - erikdesjardins:undefconst, r=oli-obk

Use undef for uninitialized bytes in constants

Fixes #83657

This generates good code when the const is fully uninit, e.g.

```rust
#[no_mangle]
pub const fn fully_uninit() -> MaybeUninit<[u8; 10]> {
    const M: MaybeUninit<[u8; 10]> = MaybeUninit::uninit();
    M
}
```
generates
```asm
fully_uninit:
	ret
```

as you would expect.

There is no improvement, however, when it's partially uninit, e.g.

```rust
pub struct PartiallyUninit {
    x: u64,
    y: MaybeUninit<[u8; 10]>
}

#[no_mangle]
pub const fn partially_uninit() -> PartiallyUninit {
    const X: PartiallyUninit = PartiallyUninit { x: 0xdeadbeefcafe, y: MaybeUninit::uninit() };
    X
}
```
generates
```asm
partially_uninit:
	mov	rax, rdi
	mov	rcx, qword ptr [rip + .L__unnamed_1+16]
	mov	qword ptr [rdi + 16], rcx
	movups	xmm0, xmmword ptr [rip + .L__unnamed_1]
	movups	xmmword ptr [rdi], xmm0
	ret

.L__unnamed_1:
	.asciz	"\376\312\357\276\255\336\000"
	.zero	16
	.size	.L__unnamed_1, 24
```
which copies a bunch of zeros in place of the undef bytes, the same as before this change.

The constant is undef where it should be, though (see the added test) so this seems to be a missed optimization in LLVM.

Author: Dylan-DPC

compiler/rustc_codegen_llvm/src/consts.rs

@@ -8,22 +8,47 @@ use crate::value::Value;
 use cstr::cstr;
 use libc::c_uint;
 use rustc_codegen_ssa::traits::*;
+use rustc_data_structures::captures::Captures;
 use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, CodegenFnAttrs};
 use rustc_middle::mir::interpret::{
-    read_target_uint, Allocation, ErrorHandled, GlobalAlloc, Pointer,
+    read_target_uint, Allocation, ErrorHandled, GlobalAlloc, InitChunk, Pointer,
 };
 use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::ty::{self, Instance, Ty};
 use rustc_middle::{bug, span_bug};
 use rustc_target::abi::{AddressSpace, Align, HasDataLayout, LayoutOf, Primitive, Scalar, Size};
+use std::ops::Range;
 use tracing::debug;
 
 pub fn const_alloc_to_llvm(cx: &CodegenCx<'ll, '_>, alloc: &Allocation) -> &'ll Value {
     let mut llvals = Vec::with_capacity(alloc.relocations().len() + 1);
     let dl = cx.data_layout();
     let pointer_size = dl.pointer_size.bytes() as usize;
 
+    // Note: this function may call `inspect_with_uninit_and_ptr_outside_interpreter`,
+    // so `range` must be within the bounds of `alloc` and not within a relocation.
+    fn chunks_of_init_and_uninit_bytes<'ll, 'a, 'b>(
+        cx: &'a CodegenCx<'ll, 'b>,
+        alloc: &'a Allocation,
+        range: Range<usize>,
+    ) -> impl Iterator<Item = &'ll Value> + Captures<'a> + Captures<'b> {
+        alloc
+            .init_mask()
+            .range_as_init_chunks(Size::from_bytes(range.start), Size::from_bytes(range.end))
+            .map(move |chunk| match chunk {
+                InitChunk::Init(range) => {
+                    let range = (range.start.bytes() as usize)..(range.end.bytes() as usize);
+                    let bytes = alloc.inspect_with_uninit_and_ptr_outside_interpreter(range);
+                    cx.const_bytes(bytes)
+                }
+                InitChunk::Uninit(range) => {
+                    let len = range.end.bytes() - range.start.bytes();
+                    cx.const_undef(cx.type_array(cx.type_i8(), len))
+                }
+            })
+    }
+
     let mut next_offset = 0;
     for &(offset, ((), alloc_id)) in alloc.relocations().iter() {
         let offset = offset.bytes();
@@ -32,12 +57,8 @@ pub fn const_alloc_to_llvm(cx: &CodegenCx<'ll, '_>, alloc: &Allocation) -> &'ll
         if offset > next_offset {
             // This `inspect` is okay since we have checked that it is not within a relocation, it
             // is within the bounds of the allocation, and it doesn't affect interpreter execution
-            // (we inspect the result after interpreter execution). Any undef byte is replaced with
-            // some arbitrary byte value.
-            //
-            // FIXME: relay undef bytes to codegen as undef const bytes
-            let bytes = alloc.inspect_with_uninit_and_ptr_outside_interpreter(next_offset..offset);
-            llvals.push(cx.const_bytes(bytes));
+            // (we inspect the result after interpreter execution).
+            llvals.extend(chunks_of_init_and_uninit_bytes(cx, alloc, next_offset..offset));
         }
         let ptr_offset = read_target_uint(
             dl.endian,
@@ -65,12 +86,8 @@ pub fn const_alloc_to_llvm(cx: &CodegenCx<'ll, '_>, alloc: &Allocation) -> &'ll
         let range = next_offset..alloc.len();
         // This `inspect` is okay since we have check that it is after all relocations, it is
         // within the bounds of the allocation, and it doesn't affect interpreter execution (we
-        // inspect the result after interpreter execution). Any undef byte is replaced with some
-        // arbitrary byte value.
-        //
-        // FIXME: relay undef bytes to codegen as undef const bytes
-        let bytes = alloc.inspect_with_uninit_and_ptr_outside_interpreter(range);
-        llvals.push(cx.const_bytes(bytes));
+        // inspect the result after interpreter execution).
+        llvals.extend(chunks_of_init_and_uninit_bytes(cx, alloc, range));
     }
 
     cx.const_struct(&llvals, true)

compiler/rustc_middle/src/mir/interpret/allocation.rs

@@ -761,20 +761,24 @@ impl InitMask {
         }
 
         // FIXME(oli-obk): optimize this for allocations larger than a block.
-        let idx = (start.bytes()..end.bytes()).map(Size::from_bytes).find(|&i| !self.get(i));
+        let idx = (start..end).find(|&i| !self.get(i));
 
         match idx {
             Some(idx) => {
-                let uninit_end = (idx.bytes()..end.bytes())
-                    .map(Size::from_bytes)
-                    .find(|&i| self.get(i))
-                    .unwrap_or(end);
+                let uninit_end = (idx..end).find(|&i| self.get(i)).unwrap_or(end);
                 Err(idx..uninit_end)
             }
             None => Ok(()),
         }
     }
 
+    /// Returns an iterator, yielding a range of byte indexes for each contiguous region
+    /// of initialized or uninitialized bytes inside the range `start..end` (end-exclusive).
+    #[inline]
+    pub fn range_as_init_chunks(&self, start: Size, end: Size) -> InitChunkIter<'_> {
+        InitChunkIter::new(self, start, end)
+    }
+
     pub fn set_range(&mut self, start: Size, end: Size, new_state: bool) {
         let len = self.len;
         if end > len {
@@ -867,6 +871,49 @@ impl InitMask {
     }
 }
 
+/// Yields [`InitChunk`]s. See [`InitMask::range_as_init_chunks`].
+pub struct InitChunkIter<'a> {
+    init_mask: &'a InitMask,
+    /// The current byte index into `init_mask`.
+    start: Size,
+    /// The end byte index into `init_mask`.
+    end: Size,
+}
+
+/// A contiguous chunk of initialized or uninitialized memory.
+pub enum InitChunk {
+    Init(Range<Size>),
+    Uninit(Range<Size>),
+}
+
+impl<'a> InitChunkIter<'a> {
+    fn new(init_mask: &'a InitMask, start: Size, end: Size) -> Self {
+        assert!(start <= end);
+        assert!(end <= init_mask.len);
+        Self { init_mask, start, end }
+    }
+}
+
+impl<'a> Iterator for InitChunkIter<'a> {
+    type Item = InitChunk;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.start >= self.end {
+            return None;
+        }
+
+        let is_init = self.init_mask.get(self.start);
+        // FIXME(oli-obk): optimize this for allocations larger than a block.
+        let end_of_chunk =
+            (self.start..self.end).find(|&i| self.init_mask.get(i) != is_init).unwrap_or(self.end);
+        let range = self.start..end_of_chunk;
+
+        self.start = end_of_chunk;
+
+        Some(if is_init { InitChunk::Init(range) } else { InitChunk::Uninit(range) })
+    }
+}
+
 #[inline]
 fn bit_index(bits: Size) -> (usize, usize) {
     let bits = bits.bytes();

compiler/rustc_middle/src/mir/interpret/mod.rs

@@ -125,7 +125,9 @@ pub use self::error::{
 
 pub use self::value::{get_slice_bytes, ConstAlloc, ConstValue, Scalar, ScalarMaybeUninit};
 
-pub use self::allocation::{Allocation, AllocationExtra, InitMask, Relocations};
+pub use self::allocation::{
+    Allocation, AllocationExtra, InitChunk, InitChunkIter, InitMask, Relocations,
+};
 
 pub use self::pointer::{Pointer, PointerArithmetic};
 

compiler/rustc_target/src/abi/mod.rs

@@ -5,6 +5,7 @@ use crate::spec::Target;
 
 use std::convert::{TryFrom, TryInto};
 use std::fmt;
+use std::iter::Step;
 use std::num::NonZeroUsize;
 use std::ops::{Add, AddAssign, Deref, Mul, Range, RangeInclusive, Sub};
 use std::str::FromStr;
@@ -433,6 +434,43 @@ impl AddAssign for Size {
     }
 }
 
+unsafe impl Step for Size {
+    #[inline]
+    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
+        u64::steps_between(&start.bytes(), &end.bytes())
+    }
+
+    #[inline]
+    fn forward_checked(start: Self, count: usize) -> Option<Self> {
+        u64::forward_checked(start.bytes(), count).map(Self::from_bytes)
+    }
+
+    #[inline]
+    fn forward(start: Self, count: usize) -> Self {
+        Self::from_bytes(u64::forward(start.bytes(), count))
+    }
+
+    #[inline]
+    unsafe fn forward_unchecked(start: Self, count: usize) -> Self {
+        Self::from_bytes(u64::forward_unchecked(start.bytes(), count))
+    }
+
+    #[inline]
+    fn backward_checked(start: Self, count: usize) -> Option<Self> {
+        u64::backward_checked(start.bytes(), count).map(Self::from_bytes)
+    }
+
+    #[inline]
+    fn backward(start: Self, count: usize) -> Self {
+        Self::from_bytes(u64::backward(start.bytes(), count))
+    }
+
+    #[inline]
+    unsafe fn backward_unchecked(start: Self, count: usize) -> Self {
+        Self::from_bytes(u64::backward_unchecked(start.bytes(), count))
+    }
+}
+
 /// Alignment of a type in bytes (always a power of two).
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Encodable, Decodable)]
 #[derive(HashStable_Generic)]

compiler/rustc_target/src/lib.rs

@@ -15,6 +15,9 @@
 #![feature(never_type)]
 #![feature(associated_type_bounds)]
 #![feature(exhaustive_patterns)]
+#![feature(step_trait)]
+#![feature(step_trait_ext)]
+#![feature(unchecked_math)]
 
 #[macro_use]
 extern crate rustc_macros;

src/test/codegen/consts.rs

@@ -43,14 +43,14 @@ pub fn inline_enum_const() -> E<i8, i16> {
 #[no_mangle]
 pub fn low_align_const() -> E<i16, [i16; 3]> {
     // Check that low_align_const and high_align_const use the same constant
-    // CHECK: memcpy.p0i8.p0i8.i{{(32|64)}}(i8* align 2 %1, i8* align 2 getelementptr inbounds (<{ [8 x i8] }>, <{ [8 x i8] }>* [[LOW_HIGH]], i32 0, i32 0, i32 0), i{{(32|64)}} 8, i1 false)
+    // CHECK: memcpy.p0i8.p0i8.i{{(32|64)}}(i8* align 2 %1, i8* align 2 getelementptr inbounds (<{ [4 x i8], [4 x i8] }>, <{ [4 x i8], [4 x i8] }>* [[LOW_HIGH]], i32 0, i32 0, i32 0), i{{(32|64)}} 8, i1 false)
     *&E::A(0)
 }
 
 // CHECK-LABEL: @high_align_const
 #[no_mangle]
 pub fn high_align_const() -> E<i16, i32> {
     // Check that low_align_const and high_align_const use the same constant
-    // CHECK: memcpy.p0i8.p0i8.i{{(32|64)}}(i8* align 4 %1, i8* align 4 getelementptr inbounds (<{ [8 x i8] }>, <{ [8 x i8] }>* [[LOW_HIGH]], i32 0, i32 0, i32 0), i{{(32|64)}} 8, i1 false)
+    // CHECK: memcpy.p0i8.p0i8.i{{(32|64)}}(i8* align 4 %1, i8* align 4 getelementptr inbounds (<{ [4 x i8], [4 x i8] }>, <{ [4 x i8], [4 x i8] }>* [[LOW_HIGH]], i32 0, i32 0, i32 0), i{{(32|64)}} 8, i1 false)
     *&E::A(0)
 }

src/test/codegen/uninit-consts.rs

@@ -0,0 +1,32 @@
+// compile-flags: -C no-prepopulate-passes
+// ignore-tidy-linelength
+
+// Check that we use undef (and not zero) for uninitialized bytes in constants.
+
+#![crate_type = "lib"]
+
+use std::mem::MaybeUninit;
+
+pub struct PartiallyUninit {
+    x: u32,
+    y: MaybeUninit<[u8; 10]>
+}
+
+// CHECK: [[FULLY_UNINIT:@[0-9]+]] = private unnamed_addr constant <{ [10 x i8] }> undef
+// CHECK: [[PARTIALLY_UNINIT:@[0-9]+]] = private unnamed_addr constant <{ [4 x i8], [12 x i8] }> <{ [4 x i8] c"\EF\BE\AD\DE", [12 x i8] undef }>, align 4
+
+// 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.p0i8.p0i8.i{{(32|64)}}(i8* align 1 %1, i8* align 1 getelementptr inbounds (<{ [10 x i8] }>, <{ [10 x i8] }>* [[FULLY_UNINIT]], i32 0, i32 0, i32 0), i{{(32|64)}} 10, i1 false)
+    M
+}
+
+// CHECK-LABEL: @partially_uninit
+#[no_mangle]
+pub const fn partially_uninit() -> PartiallyUninit {
+    const X: PartiallyUninit = PartiallyUninit { x: 0xdeadbeef, y: MaybeUninit::uninit() };
+    // CHECK: call void @llvm.memcpy.p0i8.p0i8.i{{(32|64)}}(i8* align 4 %1, i8* align 4 getelementptr inbounds (<{ [4 x i8], [12 x i8] }>, <{ [4 x i8], [12 x i8] }>* [[PARTIALLY_UNINIT]], i32 0, i32 0, i32 0), i{{(32|64)}} 16, i1 false)
+    X
+}