Use llvm.memset.p0i8.* to initialize all same-bytes arrays

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -93,23 +93,37 @@ 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
+                };
+
+                match cg_elem.val {
+                    OperandValue::Immediate(v) => {
+                        if try_init_all_same(bx, v) {
+                            return;
+                        }
                     }
+                    _ => (),
                 }
 
                 let count = self

tests/codegen/slice-init.rs

@@ -63,6 +63,32 @@ 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]
+}
+
+// FIXME: undef bytes can just be initialized with the same value as the
+// defined bytes, if the defines bytes are all the same.
+// CHECK-LABEL: @option_none_init
+#[no_mangle]
+pub fn option_none_init() -> [Option<u8>; N] {
+    // CHECK-NOT: select
+    // CHECK: br label %repeat_loop_header{{.*}}
+    // CHECK-NOT: switch
+    // CHECK: icmp
+    // CHECK-NOT: 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) {}