From 58162338a70d428a5877d68354554e6932138a9a Mon Sep 17 00:00:00 2001
From: Michal Nazarewicz <mina86@mina86.com>
Date: Thu, 9 Feb 2023 06:11:50 +0100
Subject: [PATCH] Add a safe way to create UninitSlice from slices

Introduce UninitSlice::from_slice and UninitSlice::from_uninit_slice
methods which safely create Uninit slice from provided slice of maybe
uninitialised or initialised memory.

In addition, add `From<&mut [T]>` implementations (for `T=u8` and
`T=MaybeUninit<u8>`) which do conversion from slice to UninitSlice.

Closes: https://github.com/tokio-rs/bytes/issues/552
---
 src/buf/uninit_slice.rs | 56 ++++++++++++++++++++++++++++++++++++-----
 src/bytes_mut.rs        |  2 +-
 2 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/src/buf/uninit_slice.rs b/src/buf/uninit_slice.rs
index 3161a147e..84b1d8820 100644
--- a/src/buf/uninit_slice.rs
+++ b/src/buf/uninit_slice.rs
@@ -22,10 +22,44 @@ use core::ops::{
 pub struct UninitSlice([MaybeUninit<u8>]);
 
 impl UninitSlice {
-    pub(crate) fn from_slice(slice: &mut [MaybeUninit<u8>]) -> &mut UninitSlice {
+    /// Creates a `&mut UninitSlice` wrapping slice of uninitialised memory.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::buf::UninitSlice;
+    /// use core::mem::MaybeUninit;
+    ///
+    /// let mut buffer = [MaybeUninit::uninit(); 64];
+    /// let slice = UninitSlice::from_uninit_slice(&mut buffer[..]);
+    ///
+    /// let mut vec = Vec::with_capacity(1024);
+    /// let spare: &mut UninitSlice = vec.spare_capacity_mut().into();
+    /// ```
+    #[inline]
+    pub fn from_uninit_slice(slice: &mut [MaybeUninit<u8>]) -> &mut UninitSlice {
         unsafe { &mut *(slice as *mut [MaybeUninit<u8>] as *mut UninitSlice) }
     }
 
+    fn from_uninit_slice_ref(slice: &[MaybeUninit<u8>]) -> &UninitSlice {
+        unsafe { &*(slice as *const [MaybeUninit<u8>] as *const UninitSlice) }
+    }
+
+    /// Creates a `&mut UninitSlice` wrapping slice of initialised memory.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::buf::UninitSlice;
+    ///
+    /// let mut buffer = [0u8; 64];
+    /// let slice = UninitSlice::from_slice(&mut buffer[..]);
+    /// ```
+    #[inline]
+    pub fn from_slice(slice: &mut [u8]) -> &mut UninitSlice {
+        unsafe { &mut *(slice as *mut [u8] as *mut [MaybeUninit<u8>] as *mut UninitSlice) }
+    }
+
     /// Create a `&mut UninitSlice` from a pointer and a length.
     ///
     /// # Safety
@@ -48,7 +82,7 @@ impl UninitSlice {
     pub unsafe fn from_raw_parts_mut<'a>(ptr: *mut u8, len: usize) -> &'a mut UninitSlice {
         let maybe_init: &mut [MaybeUninit<u8>] =
             core::slice::from_raw_parts_mut(ptr as *mut _, len);
-        Self::from_slice(maybe_init)
+        Self::from_uninit_slice(maybe_init)
     }
 
     /// Write a single byte at the specified offset.
@@ -179,6 +213,18 @@ impl fmt::Debug for UninitSlice {
     }
 }
 
+impl<'a> From<&'a mut [u8]> for &'a mut UninitSlice {
+    fn from(slice: &'a mut [u8]) -> Self {
+        UninitSlice::from_slice(slice)
+    }
+}
+
+impl<'a> From<&'a mut [MaybeUninit<u8>]> for &'a mut UninitSlice {
+    fn from(slice: &'a mut [MaybeUninit<u8>]) -> Self {
+        UninitSlice::from_uninit_slice(slice)
+    }
+}
+
 macro_rules! impl_index {
     ($($t:ty),*) => {
         $(
@@ -187,16 +233,14 @@ macro_rules! impl_index {
 
                 #[inline]
                 fn index(&self, index: $t) -> &UninitSlice {
-                    let maybe_uninit: &[MaybeUninit<u8>] = &self.0[index];
-                    unsafe { &*(maybe_uninit as *const [MaybeUninit<u8>] as *const UninitSlice) }
+                    UninitSlice::from_uninit_slice_ref(&self.0[index])
                 }
             }
 
             impl IndexMut<$t> for UninitSlice {
                 #[inline]
                 fn index_mut(&mut self, index: $t) -> &mut UninitSlice {
-                    let maybe_uninit: &mut [MaybeUninit<u8>] = &mut self.0[index];
-                    unsafe { &mut *(maybe_uninit as *mut [MaybeUninit<u8>] as *mut UninitSlice) }
+                    UninitSlice::from_uninit_slice(&mut self.0[index])
                 }
             }
         )*
diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs
index 450b93279..c5c2e52fc 100644
--- a/src/bytes_mut.rs
+++ b/src/bytes_mut.rs
@@ -1102,7 +1102,7 @@ unsafe impl BufMut for BytesMut {
         if self.capacity() == self.len() {
             self.reserve(64);
         }
-        UninitSlice::from_slice(self.spare_capacity_mut())
+        self.spare_capacity_mut().into()
     }
 
     // Specialize these methods so they can skip checking `remaining_mut`