Make intrinsics::write_bytes const

Author: lilasta

compiler/rustc_const_eval/src/interpret/intrinsics.rs

@@ -322,6 +322,9 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             sym::copy => {
                 self.copy_intrinsic(&args[0], &args[1], &args[2], /*nonoverlapping*/ false)?;
             }
+            sym::write_bytes => {
+                self.write_bytes_intrinsic(&args[0], &args[1], &args[2])?;
+            }
             sym::offset => {
                 let ptr = self.read_pointer(&args[0])?;
                 let offset_count = self.read_scalar(&args[1])?.to_machine_isize(self)?;
@@ -543,6 +546,27 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
         self.memory.copy(src, align, dst, align, size, nonoverlapping)
     }
 
+    pub(crate) fn write_bytes_intrinsic(
+        &mut self,
+        dst: &OpTy<'tcx, <M as Machine<'mir, 'tcx>>::PointerTag>,
+        byte: &OpTy<'tcx, <M as Machine<'mir, 'tcx>>::PointerTag>,
+        count: &OpTy<'tcx, <M as Machine<'mir, 'tcx>>::PointerTag>,
+    ) -> InterpResult<'tcx> {
+        let layout = self.layout_of(dst.layout.ty.builtin_deref(true).unwrap().ty)?;
+
+        let dst = self.read_pointer(&dst)?;
+        let byte = self.read_scalar(&byte)?.to_u8()?;
+        let count = self.read_scalar(&count)?.to_machine_usize(self)?;
+
+        let len = layout
+            .size
+            .checked_mul(count, self)
+            .ok_or_else(|| err_ub_format!("overflow computing total size of `write_bytes`"))?;
+
+        let bytes = std::iter::repeat(byte).take(len.bytes_usize());
+        self.memory.write_bytes(dst, bytes)
+    }
+
     pub(crate) fn raw_eq_intrinsic(
         &mut self,
         lhs: &OpTy<'tcx, <M as Machine<'mir, 'tcx>>::PointerTag>,

library/core/src/intrinsics.rs

@@ -2242,13 +2242,29 @@ pub const unsafe fn copy<T>(src: *const T, dst: *mut T, count: usize) {
 /// assert_eq!(*v, 42);
 /// ```
 #[stable(feature = "rust1", since = "1.0.0")]
+#[rustc_const_unstable(feature = "const_ptr_write", issue = "86302")]
 #[inline]
-pub unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize) {
+pub const unsafe fn write_bytes<T>(dst: *mut T, val: u8, count: usize) {
     extern "rust-intrinsic" {
+        #[rustc_const_unstable(feature = "const_ptr_write", issue = "86302")]
         fn write_bytes<T>(dst: *mut T, val: u8, count: usize);
     }
 
-    debug_assert!(is_aligned_and_not_null(dst), "attempt to write to unaligned or null pointer");
+    #[cfg(debug_assertions)]
+    fn runtime_check<T>(ptr: *mut T) {
+        debug_assert!(
+            is_aligned_and_not_null(ptr),
+            "attempt to write to unaligned or null pointer"
+        );
+    }
+    #[cfg(debug_assertions)]
+    const fn compiletime_check<T>(_ptr: *mut T) {}
+    #[cfg(debug_assertions)]
+    // SAFETY: runtime debug-assertions are a best-effort basis; it's fine to
+    // not do them during compile time
+    unsafe {
+        const_eval_select((dst,), compiletime_check, runtime_check);
+    }
 
     // SAFETY: the safety contract for `write_bytes` must be upheld by the caller.
     unsafe { write_bytes(dst, val, count) }

library/core/tests/intrinsics.rs

@@ -35,3 +35,33 @@ fn test_assume_can_be_in_const_contexts() {
     let rs = unsafe { foo(42, 97) };
     assert_eq!(rs, 0);
 }
+
+#[test]
+#[cfg(not(bootstrap))]
+const fn test_write_bytes_in_const_contexts() {
+    use core::intrinsics::write_bytes;
+
+    const TEST: [u32; 3] = {
+        let mut arr = [1u32, 2, 3];
+        unsafe {
+            write_bytes(arr.as_mut_ptr(), 0, 2);
+        }
+        arr
+    };
+
+    assert!(TEST[0] == 0);
+    assert!(TEST[1] == 0);
+    assert!(TEST[2] == 3);
+
+    const TEST2: [u32; 3] = {
+        let mut arr = [1u32, 2, 3];
+        unsafe {
+            write_bytes(arr.as_mut_ptr(), 1, 2);
+        }
+        arr
+    };
+
+    assert!(TEST2[0] == 16843009);
+    assert!(TEST2[1] == 16843009);
+    assert!(TEST2[2] == 3);
+}