Skip to content

Commit

Permalink
Add a safe way to create UninitSlice from slices
Browse files Browse the repository at this point in the history
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: tokio-rs#552
  • Loading branch information
mina86 committed Feb 9, 2023
1 parent 74b04c7 commit 5816233
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 7 deletions.
56 changes: 50 additions & 6 deletions src/buf/uninit_slice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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),*) => {
$(
Expand All @@ -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])
}
}
)*
Expand Down
2 changes: 1 addition & 1 deletion src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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`
Expand Down

0 comments on commit 5816233

Please sign in to comment.