specialize slice::fill to use memset

compiler/rustc_codegen_cranelift/example/mini_core.rs

@@ -666,7 +666,7 @@ pub mod intrinsics {
     #[rustc_intrinsic]
     pub fn bswap<T>(x: T) -> T;
     #[rustc_intrinsic]
-    pub unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize);
+    pub unsafe fn write_bytes<T, B>(dst: *mut T, val: B, count: usize);
     #[rustc_intrinsic]
     pub unsafe fn unreachable() -> !;
 }

compiler/rustc_codegen_cranelift/example/mini_core_hello_world.rs

@@ -125,7 +125,7 @@ static NUM_REF: &'static u8 = unsafe { &*&raw const NUM };
 
 unsafe fn zeroed<T>() -> T {
     let mut uninit = MaybeUninit { uninit: () };
-    intrinsics::write_bytes(&mut uninit.value.value as *mut T, 0, 1);
+    intrinsics::write_bytes(&mut uninit.value.value as *mut T, 0u8, 1);
     uninit.value.value
 }
 

compiler/rustc_codegen_cranelift/src/intrinsics/mod.rs

@@ -685,6 +685,12 @@ fn codegen_regular_intrinsic_call<'tcx>(
 
         sym::write_bytes | sym::volatile_set_memory => {
             intrinsic_args!(fx, args => (dst, val, count); intrinsic);
+            if val.layout().size.bytes() != 1 {
+                // incorrect sizes can be encountered on dead branches
+                fx.bcx.ins().trap(TrapCode::user(1).unwrap());
+                return Ok(());
+            };
+
             let val = val.load_scalar(fx);
             let count = count.load_scalar(fx);
 

compiler/rustc_codegen_gcc/example/mini_core.rs

@@ -677,7 +677,7 @@ pub mod intrinsics {
     #[rustc_intrinsic]
     pub fn bswap<T>(x: T) -> T;
     #[rustc_intrinsic]
-    pub unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize);
+    pub unsafe fn write_bytes<T, B>(dst: *mut T, val: B, count: usize);
     #[rustc_intrinsic]
     pub unsafe fn unreachable() -> !;
 }

compiler/rustc_codegen_gcc/example/mini_core_hello_world.rs

@@ -127,7 +127,7 @@ impl<T: ?Sized, U: ?Sized> CoerceUnsized<Unique<U>> for Unique<T> where T: Unsiz
 
 unsafe fn zeroed<T>() -> T {
     let mut uninit = MaybeUninit { uninit: () };
-    intrinsics::write_bytes(&mut uninit.value.value as *mut T, 0, 1);
+    intrinsics::write_bytes(&mut uninit.value.value as *mut T, 0u8, 1);
     uninit.value.value
 }
 

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -198,12 +198,18 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 return Ok(());
             }
             sym::write_bytes => {
+                // invalid types may be encountered on dead branches after a size check.
+                if args[1].layout.size.bytes() != 1 {
+                    bx.unreachable_nonterminator();
+                    return Ok(());
+                }
+                let val = bx.from_immediate(args[1].immediate());
                 memset_intrinsic(
                     bx,
                     false,
                     fn_args.type_at(0),
                     args[0].immediate(),
-                    args[1].immediate(),
+                    val,
                     args[2].immediate(),
                 );
                 return Ok(());

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -4,6 +4,7 @@
 
 use std::assert_matches::assert_matches;
 
+use either::Either;
 use rustc_abi::{FieldIdx, HasDataLayout, Size};
 use rustc_apfloat::ieee::{Double, Half, Quad, Single};
 use rustc_middle::mir::interpret::{CTFE_ALLOC_SALT, read_target_uint, write_target_uint};
@@ -866,18 +867,43 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
         count: &OpTy<'tcx, <M as Machine<'tcx>>::Provenance>,
         name: &'static str,
     ) -> InterpResult<'tcx> {
-        let layout = self.layout_of(dst.layout.ty.builtin_deref(true).unwrap())?;
+        let dst_layout = self.layout_of(dst.layout.ty.builtin_deref(true).unwrap())?;
+        let src_layout = self.layout_of(byte.layout.ty)?;
+
+        if src_layout.size.bytes_usize() != 1 {
+            throw_ub_custom!(
+                fluent::const_eval_scalar_size_mismatch,
+                target_size = 1,
+                data_size = src_layout.size.bytes(),
+            );
+        }
 
         let dst = self.read_pointer(dst)?;
-        let byte = self.read_scalar(byte)?.to_u8()?;
+
         let count = self.read_target_usize(count)?;
 
         // `checked_mul` enforces a too small bound (the correct one would probably be target_isize_max),
         // but no actual allocation can be big enough for the difference to be noticeable.
         let len = self
-            .compute_size_in_bytes(layout.size, count)
+            .compute_size_in_bytes(dst_layout.size, count)
             .ok_or_else(|| err_ub_custom!(fluent::const_eval_size_overflow, name = name))?;
 
+        let byte = match self.read_immediate_raw(byte)? {
+            Either::Left(src_place) => {
+                // val is not an immediate, possibly uninit.
+                self.mem_copy_repeatedly(
+                    src_place.ptr(),
+                    dst,
+                    Size::from_bytes(1),
+                    len.bytes(),
+                    /* nonoverlapping: */ false,
+                )?;
+                return interp_ok(());
+            }
+            Either::Right(imm) => imm,
+        };
+
+        let byte = byte.to_scalar().to_u8()?;
         let bytes = std::iter::repeat(byte).take(len.bytes_usize());
         self.write_bytes_ptr(dst, bytes)
     }

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -341,7 +341,10 @@ pub(crate) fn check_intrinsic_type(
             let byte_ptr = Ty::new_imm_ptr(tcx, tcx.types.u8);
             (0, 0, vec![byte_ptr, byte_ptr, tcx.types.usize], tcx.types.i32)
         }
-        sym::write_bytes | sym::volatile_set_memory => (
+        sym::write_bytes => {
+            (2, 0, vec![Ty::new_mut_ptr(tcx, param(0)), param(1), tcx.types.usize], tcx.types.unit)
+        }
+        sym::volatile_set_memory => (
             1,
             0,
             vec![Ty::new_mut_ptr(tcx, param(0)), tcx.types.u8, tcx.types.usize],

library/core/src/intrinsics/mod.rs

@@ -2882,6 +2882,8 @@ pub const unsafe fn copy_nonoverlapping<T>(src: *const T, dst: *mut T, count: us
 pub const unsafe fn copy<T>(src: *const T, dst: *mut T, count: usize);
 
 /// This is an accidentally-stable alias to [`ptr::write_bytes`]; use that instead.
+///
+/// `val` must be 1 byte wide, it is allowed to be uninit.
 // Note (intentionally not in the doc comment): `ptr::write_bytes` adds some extra
 // debug assertions; if you are writing compiler tests or code inside the standard library
 // that wants to avoid those debug assertions, directly call this intrinsic instead.
@@ -2890,7 +2892,7 @@ pub const unsafe fn copy<T>(src: *const T, dst: *mut T, count: usize);
 #[rustc_const_stable(feature = "const_intrinsic_copy", since = "1.83.0")]
 #[rustc_nounwind]
 #[rustc_intrinsic]
-pub const unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize);
+pub const unsafe fn write_bytes<T, B>(dst: *mut T, val: B, count: usize);
 
 /// Returns the minimum (IEEE 754-2008 minNum) of two `f16` values.
 ///

library/core/src/slice/specialize.rs

@@ -15,9 +15,43 @@ impl<T: Clone> SpecFill<T> for [T] {
 }
 
 impl<T: Copy> SpecFill<T> for [T] {
-    fn spec_fill(&mut self, value: T) {
+    default fn spec_fill(&mut self, value: T) {
+        if size_of::<T>() == 1 {
+            // SAFETY: The pointer is derived from a reference, so it's writable.
+            // And we checked that T is 1 byte wide.
+            unsafe {
+                // use the intrinsic since it allows any T as long as it's 1 byte wide
+                crate::intrinsics::write_bytes(self.as_mut_ptr(), value, self.len());
+            }
+            return;
+        }
         for item in self.iter_mut() {
             *item = value;
         }
     }
 }
+
+macro spec_fill_int {
+    ($($type:ty)*) => {$(
+        impl SpecFill<$type> for [$type] {
+            #[inline]
+            fn spec_fill(&mut self, value: $type) {
+                if crate::intrinsics::is_val_statically_known(value) {
+                    let bytes = value.to_ne_bytes();
+                    if value == <$type>::from_ne_bytes([bytes[0]; size_of::<$type>()]) {
+                        // SAFETY: The pointer is derived from a reference, so it's writable.
+                        unsafe {
+                            crate::intrinsics::write_bytes(self.as_mut_ptr(), bytes[0], self.len());
+                        }
+                        return;
+                    }
+                }
+                for item in self.iter_mut() {
+                    *item = value;
+                }
+            }
+        }
+    )*}
+}
+
+spec_fill_int! { u16 i16 u32 i32 u64 i64 u128 i128 usize isize }

library/coretests/tests/intrinsics.rs

@@ -1,5 +1,7 @@
 use core::any::TypeId;
+use core::hint::black_box;
 use core::intrinsics::assume;
+use core::mem::MaybeUninit;
 
 #[test]
 fn test_typeid_sized_types() {
@@ -43,7 +45,7 @@ const fn test_write_bytes_in_const_contexts() {
     const TEST: [u32; 3] = {
         let mut arr = [1u32, 2, 3];
         unsafe {
-            write_bytes(arr.as_mut_ptr(), 0, 2);
+            write_bytes(arr.as_mut_ptr(), 0u8, 2);
         }
         arr
     };
@@ -55,14 +57,25 @@ const fn test_write_bytes_in_const_contexts() {
     const TEST2: [u32; 3] = {
         let mut arr = [1u32, 2, 3];
         unsafe {
-            write_bytes(arr.as_mut_ptr(), 1, 2);
+            write_bytes(arr.as_mut_ptr(), 1u8, 2);
         }
         arr
     };
 
     assert!(TEST2[0] == 16843009);
     assert!(TEST2[1] == 16843009);
     assert!(TEST2[2] == 3);
+
+    const TEST3: [MaybeUninit<u32>; 2] = {
+        let mut arr: [MaybeUninit<u32>; 2] = [MaybeUninit::uninit(), MaybeUninit::new(1)];
+        unsafe {
+            write_bytes(arr.as_mut_ptr(), MaybeUninit::<u8>::uninit(), 2);
+        }
+        arr
+    };
+
+    // can't do much with uninitialized values, just make sure it compiles
+    black_box(TEST3);
 }
 
 #[test]

tests/codegen-llvm/lib-optimizations/slice_fill.rs

@@ -0,0 +1,28 @@
+//@ compile-flags: -Copt-level=3
+#![crate_type = "lib"]
+
+use std::mem::MaybeUninit;
+
+// CHECK-LABEL: @slice_fill_pass_undef
+#[no_mangle]
+pub fn slice_fill_pass_undef(s: &mut [MaybeUninit<u8>], v: MaybeUninit<u8>) {
+    // CHECK: tail call void @llvm.memset.{{.*}}(ptr nonnull align 1 %s.0, i8 %v, {{.*}} %s.1, i1 false)
+    // CHECK: ret
+    s.fill(v);
+}
+
+// CHECK-LABEL: @slice_fill_uninit
+#[no_mangle]
+pub fn slice_fill_uninit(s: &mut [MaybeUninit<u8>]) {
+    // CHECK-NOT: call
+    // CHECK: ret void
+    s.fill(MaybeUninit::uninit());
+}
+
+// CHECK-LABEL: @slice_wide_memset
+#[no_mangle]
+pub fn slice_wide_memset(s: &mut [u16]) {
+    // CHECK: tail call void @llvm.memset.{{.*}}(ptr nonnull align 2 %s.0, i8 -1
+    // CHECK: ret
+    s.fill(0xFFFF);
+}