generalize intrinsics::write_bytes to take any 1-byte wide type as value

This lets us use it for any T since T may contain uninitialized bytes.

Author: the8472

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,40 @@ 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.
+                for i in 0..len.bytes_usize() {
+                    let ptr = dst.wrapping_offset(Size::from_bytes(i), &self.tcx);
+                    self.mem_copy(src_place.ptr(), ptr, Size::from_bytes(1), true)?;
+                }
+                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/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]