Auto merge of rust-lang#135335 - oli-obk:push-zxwssomxxtnq, r=<try>

Treat undef bytes as equal to any other byte

Basically since `undef` can be any byte, it can also be the byte(s) that are in the non-undef parts of a value. So we can just treat the `undef` at not being there and only look at the initialized bytes and memset over them

fixes rust-lang#104290

based on rust-lang#135258

Author: bors

compiler/rustc_codegen_llvm/src/common.rs

@@ -126,6 +126,10 @@ impl<'ll, 'tcx> ConstCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
         unsafe { llvm::LLVMGetUndef(t) }
     }
 
+    fn is_undef(&self, v: &'ll Value) -> bool {
+        unsafe { llvm::LLVMIsUndef(v) == True }
+    }
+
     fn const_poison(&self, t: &'ll Type) -> &'ll Value {
         unsafe { llvm::LLVMGetPoison(t) }
     }

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -917,6 +917,7 @@ unsafe extern "C" {
     pub fn LLVMMetadataTypeInContext(C: &Context) -> &Type;
 
     // Operations on all values
+    pub fn LLVMIsUndef(Val: &Value) -> Bool;
     pub fn LLVMTypeOf(Val: &Value) -> &Type;
     pub fn LLVMGetValueName2(Val: &Value, Length: *mut size_t) -> *const c_char;
     pub fn LLVMSetValueName2(Val: &Value, Name: *const c_char, NameLen: size_t);

compiler/rustc_codegen_ssa/src/mir/operand.rs

@@ -205,13 +205,22 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
         assert!(alloc_align >= layout.align.abi);
 
         let read_scalar = |start, size, s: abi::Scalar, ty| {
+            let range = alloc_range(start, size);
             match alloc.0.read_scalar(
                 bx,
-                alloc_range(start, size),
+                range,
                 /*read_provenance*/ matches!(s.primitive(), abi::Primitive::Pointer(_)),
             ) {
-                Ok(val) => bx.scalar_to_backend(val, s, ty),
-                Err(_) => bx.const_poison(ty),
+                Ok(val) => Some(bx.scalar_to_backend(val, s, ty)),
+                Err(_) => {
+                    if alloc.0.provenance().range_empty(range, &bx.tcx())
+                        && alloc.0.init_mask().is_range_initialized(range).unwrap_err() == range
+                    {
+                        Some(bx.const_undef(ty))
+                    } else {
+                        None
+                    }
+                }
             }
         };
 
@@ -222,16 +231,14 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
         // check that walks over the type of `mplace` to make sure it is truly correct to treat this
         // like a `Scalar` (or `ScalarPair`).
         match layout.backend_repr {
-            BackendRepr::Scalar(s @ abi::Scalar::Initialized { .. }) => {
+            BackendRepr::Scalar(s) => {
                 let size = s.size(bx);
                 assert_eq!(size, layout.size, "abi::Scalar size does not match layout size");
-                let val = read_scalar(offset, size, s, bx.immediate_backend_type(layout));
-                OperandRef { val: OperandValue::Immediate(val), layout }
+                if let Some(val) = read_scalar(offset, size, s, bx.immediate_backend_type(layout)) {
+                    return OperandRef { val: OperandValue::Immediate(val), layout };
+                }
             }
-            BackendRepr::ScalarPair(
-                a @ abi::Scalar::Initialized { .. },
-                b @ abi::Scalar::Initialized { .. },
-            ) => {
+            BackendRepr::ScalarPair(a, b) => {
                 let (a_size, b_size) = (a.size(bx), b.size(bx));
                 let b_offset = (offset + a_size).align_to(b.align(bx).abi);
                 assert!(b_offset.bytes() > 0);
@@ -247,20 +254,21 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                     b,
                     bx.scalar_pair_element_backend_type(layout, 1, true),
                 );
-                OperandRef { val: OperandValue::Pair(a_val, b_val), layout }
-            }
-            _ if layout.is_zst() => OperandRef::zero_sized(layout),
-            _ => {
-                // Neither a scalar nor scalar pair. Load from a place
-                // FIXME: should we cache `const_data_from_alloc` to avoid repeating this for the
-                // same `ConstAllocation`?
-                let init = bx.const_data_from_alloc(alloc);
-                let base_addr = bx.static_addr_of(init, alloc_align, None);
-
-                let llval = bx.const_ptr_byte_offset(base_addr, offset);
-                bx.load_operand(PlaceRef::new_sized(llval, layout))
+                if let (Some(a_val), Some(b_val)) = (a_val, b_val) {
+                    return OperandRef { val: OperandValue::Pair(a_val, b_val), layout };
+                }
             }
+            _ if layout.is_zst() => return OperandRef::zero_sized(layout),
+            _ => {}
         }
+        // Neither a scalar nor scalar pair. Load from a place
+        // FIXME: should we cache `const_data_from_alloc` to avoid repeating this for the
+        // same `ConstAllocation`?
+        let init = bx.const_data_from_alloc(alloc);
+        let base_addr = bx.static_addr_of(init, alloc_align, None);
+
+        let llval = bx.const_ptr_byte_offset(base_addr, offset);
+        bx.load_operand(PlaceRef::new_sized(llval, layout))
     }
 
     /// Asserts that this operand refers to a scalar and returns

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -8,7 +8,7 @@ use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_middle::{bug, mir, span_bug};
 use rustc_session::config::OptLevel;
 use rustc_span::{DUMMY_SP, Span};
-use tracing::{debug, instrument};
+use tracing::{debug, instrument, trace};
 
 use super::FunctionCx;
 use super::operand::{OperandRef, OperandValue};
@@ -93,23 +93,57 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     return;
                 }
 
-                if let OperandValue::Immediate(v) = cg_elem.val {
+                let try_init_all_same = |bx: &mut Bx, v| {
                     let start = dest.val.llval;
                     let size = bx.const_usize(dest.layout.size.bytes());
 
-                    // Use llvm.memset.p0i8.* to initialize all zero arrays
-                    if bx.cx().const_to_opt_u128(v, false) == Some(0) {
-                        let fill = bx.cx().const_u8(0);
-                        bx.memset(start, fill, size, dest.val.align, MemFlags::empty());
-                        return;
+                    // Use llvm.memset.p0i8.* to initialize all same byte arrays
+                    if let Some(int) = bx.cx().const_to_opt_u128(v, false) {
+                        let bytes = &int.to_le_bytes()[..cg_elem.layout.size.bytes_usize()];
+                        let first = bytes[0];
+                        if bytes[1..].iter().all(|&b| b == first) {
+                            let fill = bx.cx().const_u8(first);
+                            bx.memset(start, fill, size, dest.val.align, MemFlags::empty());
+                            return true;
+                        }
                     }
 
                     // Use llvm.memset.p0i8.* to initialize byte arrays
                     let v = bx.from_immediate(v);
                     if bx.cx().val_ty(v) == bx.cx().type_i8() {
                         bx.memset(start, v, size, dest.val.align, MemFlags::empty());
-                        return;
+                        return true;
+                    }
+                    false
+                };
+
+                trace!(?cg_elem.val);
+                match cg_elem.val {
+                    OperandValue::Immediate(v) => {
+                        if try_init_all_same(bx, v) {
+                            return;
+                        }
+                    }
+                    OperandValue::Pair(a, b) => {
+                        let a_is_undef = bx.cx().is_undef(a);
+                        match (a_is_undef, bx.cx().is_undef(b)) {
+                            // Can happen for uninit unions
+                            (true, true) => {
+                                // FIXME: can we produce better output here?
+                            }
+                            (false, true) | (true, false) => {
+                                let val = if a_is_undef { b } else { b };
+                                if try_init_all_same(bx, val) {
+                                    return;
+                                }
+                            }
+                            (false, false) => {
+                                // FIXME: if both are the same value, use try_init_all_same
+                            }
+                        }
                     }
+                    OperandValue::ZeroSized => unreachable!("checked above"),
+                    OperandValue::Ref(..) => {}
                 }
 
                 let count = self

compiler/rustc_codegen_ssa/src/traits/consts.rs

@@ -9,6 +9,7 @@ pub trait ConstCodegenMethods<'tcx>: BackendTypes {
     /// Generate an uninitialized value (matching uninitialized memory in MIR).
     /// Whether memory is initialized or not is tracked byte-for-byte.
     fn const_undef(&self, t: Self::Type) -> Self::Value;
+    fn is_undef(&self, v: Self::Value) -> bool;
     /// Generate a fake value. Poison always affects the entire value, even if just a single byte is
     /// poison. This can only be used in codepaths that are already UB, i.e., UB-free Rust code
     /// (including code that e.g. copies uninit memory with `MaybeUninit`) can never encounter a

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

@@ -222,7 +222,7 @@ impl AllocError {
 }
 
 /// The information that makes up a memory access: offset and size.
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, PartialEq)]
 pub struct AllocRange {
     pub start: Size,
     pub size: Size,

tests/codegen/overaligned-constant.rs

@@ -17,8 +17,6 @@ pub fn overaligned_constant() {
     // CHECK-LABEL: @overaligned_constant
     // CHECK: [[full:%_.*]] = alloca [32 x i8], align 8
     // CHECK: call void @llvm.memcpy.p0.p0.i64(ptr align 8 [[full]], ptr align 8 @0, i64 32, i1 false)
-    // CHECK: %b.0 = load i32, ptr @0, align 4
-    // CHECK: %b.1 = load i32, ptr getelementptr inbounds ({{.*}}), align 4
     let mut s = S(1);
 
     s.0 = 3;

tests/codegen/slice-init.rs

@@ -63,6 +63,30 @@ pub fn nonzero_integer_array() {
     opaque(&x);
 }
 
+const N: usize = 100;
+
+// CHECK-LABEL: @u16_init_one_bytes
+#[no_mangle]
+pub fn u16_init_one_bytes() -> [u16; N] {
+    // CHECK-NOT: select
+    // CHECK-NOT: br
+    // CHECK-NOT: switch
+    // CHECK-NOT: icmp
+    // CHECK: call void @llvm.memset.p0
+    [const { u16::from_be_bytes([1, 1]) }; N]
+}
+
+// CHECK-LABEL: @option_none_init
+#[no_mangle]
+pub fn option_none_init() -> [Option<u8>; N] {
+    // CHECK-NOT: select
+    // CHECK-NOT: br
+    // CHECK-NOT: switch
+    // CHECK-NOT: icmp
+    // CHECK: call void @llvm.memset.p0
+    [None; N]
+}
+
 // Use an opaque function to prevent rustc from removing useless drops.
 #[inline(never)]
 pub fn opaque(_: impl Sized) {}