From c818d8b5abf8f89b56fbdf581593beacb798d43f Mon Sep 17 00:00:00 2001
From: Michal Nazarewicz <mina86@mina86.com>
Date: Thu, 9 Feb 2023 01:50:44 +0100
Subject: [PATCH] Implement BufMut for `&mut [MaybeUninit<u8>]`

This will allow using BufMut without having to initialise memory
first.  In particular, this will allow simple integration with tokio:

    #![feature(maybe_uninit_slice)]

    use tokio::io::AsyncReadExt;

    let mut tokio_file = /* ... */;
    let mut buf = [MaybeUninit::uninit(); 4096];
    let len = tokio_file.read_buf(&mut buf);
    // SAFETY: read above initialised first len bytes of the buffer.
    let buf = unsafe {
        MaybeUninit::slice_assume_init_mut(&buf[..len])
    };
---
 src/buf/buf_mut.rs | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/src/buf/buf_mut.rs b/src/buf/buf_mut.rs
index 685fcc76b..b2f2d8cc3 100644
--- a/src/buf/buf_mut.rs
+++ b/src/buf/buf_mut.rs
@@ -1419,6 +1419,41 @@ unsafe impl BufMut for &mut [u8] {
     }
 }
 
+unsafe impl BufMut for &mut [core::mem::MaybeUninit<u8>] {
+    #[inline]
+    fn remaining_mut(&self) -> usize {
+        self.len()
+    }
+
+    #[inline]
+    fn chunk_mut(&mut self) -> &mut UninitSlice {
+        UninitSlice::from_slice(self)
+    }
+
+    #[inline]
+    unsafe fn advance_mut(&mut self, cnt: usize) {
+        // Lifetime dance taken from `impl Write for &mut [u8]`.
+        let (_, b) = core::mem::replace(self, &mut []).split_at_mut(cnt);
+        *self = b;
+    }
+
+    #[inline]
+    fn put_slice(&mut self, src: &[u8]) {
+        self.chunk_mut().copy_from_slice(src);
+        unsafe {
+            self.advance_mut(src.len());
+        }
+    }
+
+    fn put_bytes(&mut self, val: u8, cnt: usize) {
+        assert!(self.remaining_mut() >= cnt);
+        unsafe {
+            ptr::write_bytes(self.as_mut_ptr() as *mut u8, val, cnt);
+            self.advance_mut(cnt);
+        }
+    }
+}
+
 unsafe impl BufMut for Vec<u8> {
     #[inline]
     fn remaining_mut(&self) -> usize {