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

2 files changed, 12 insertions, 10 deletions

diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
index e36e0e18d9c..31793641d75 100644
--- a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -97,11 +97,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     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 true;
+                    // 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
diff --git a/tests/codegen/slice-init.rs b/tests/codegen/slice-init.rs
index a01b432add6..1c2dd3e8875 100644
--- a/tests/codegen/slice-init.rs
+++ b/tests/codegen/slice-init.rs
@@ -65,16 +65,14 @@ pub fn nonzero_integer_array() {
 
 const N: usize = 100;
 
-// FIXME: The two bytes of the u16 are the same, so we should
-// just use memset, too.
 // CHECK-LABEL: @u16_init_one_bytes
 #[no_mangle]
 pub fn u16_init_one_bytes() -> [u16; N] {
     // CHECK-NOT: select
-    // CHECK: br
+    // CHECK-NOT: br
     // CHECK-NOT: switch
-    // CHECK: icmp
-    // CHECK-NOT: call void @llvm.memset.p0
+    // CHECK-NOT: icmp
+    // CHECK: call void @llvm.memset.p0
     [const { u16::from_be_bytes([1, 1]) }; N]
 }