Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion src/buf/uninit_slice.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,10 @@ use core::ops::{
pub struct UninitSlice([MaybeUninit<u8>]);

impl UninitSlice {
pub(crate) fn from_slice(slice: &mut [MaybeUninit<u8>]) -> &mut UninitSlice {
unsafe { &mut *(slice as *mut [MaybeUninit<u8>] as *mut UninitSlice) }
}

/// Create a `&mut UninitSlice` from a pointer and a length.
///
/// # Safety
Expand All @@ -44,7 +48,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);
&mut *(maybe_init as *mut [MaybeUninit<u8>] as *mut UninitSlice)
Self::from_slice(maybe_init)
}

/// Write a single byte at the specified offset.
Expand Down
43 changes: 36 additions & 7 deletions src/bytes_mut.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use core::iter::{FromIterator, Iterator};
use core::mem::{self, ManuallyDrop};
use core::mem::{self, ManuallyDrop, MaybeUninit};
use core::ops::{Deref, DerefMut};
use core::ptr::{self, NonNull};
use core::{cmp, fmt, hash, isize, slice, usize};
Expand Down Expand Up @@ -766,11 +766,11 @@ impl BytesMut {
self.reserve(cnt);

unsafe {
let dst = self.uninit_slice();
let dst = self.spare_capacity_mut();
// Reserved above
debug_assert!(dst.len() >= cnt);

ptr::copy_nonoverlapping(extend.as_ptr(), dst.as_mut_ptr(), cnt);
ptr::copy_nonoverlapping(extend.as_ptr(), dst.as_mut_ptr().cast(), cnt);
}

unsafe {
Expand Down Expand Up @@ -992,13 +992,42 @@ impl BytesMut {
self.data = invalid_ptr((pos << VEC_POS_OFFSET) | (prev & NOT_VEC_POS_MASK));
}

/// Returns the remaining spare capacity of the buffer as a slice of `MaybeUninit<u8>`.
///
/// The returned slice can be used to fill the buffer with data (e.g. by
/// reading from a file) before marking the data as initialized using the
/// [`set_len`] method.
///
/// [`set_len`]: BytesMut::set_len
///
/// # Examples
///
/// ```
/// use bytes::BytesMut;
///
/// // Allocate buffer big enough for 10 bytes.
/// let mut buf = BytesMut::with_capacity(10);
///
/// // Fill in the first 3 elements.
/// let uninit = buf.spare_capacity_mut();
/// uninit[0].write(0);
/// uninit[1].write(1);
/// uninit[2].write(2);
///
/// // Mark the first 3 bytes of the buffer as being initialized.
/// unsafe {
/// buf.set_len(3);
/// }
///
/// assert_eq!(&buf[..], &[0, 1, 2]);
/// ```
#[inline]
fn uninit_slice(&mut self) -> &mut UninitSlice {
pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<u8>] {
unsafe {
let ptr = self.ptr.as_ptr().add(self.len);
let len = self.cap - self.len;

UninitSlice::from_raw_parts_mut(ptr, len)
slice::from_raw_parts_mut(ptr.cast(), len)
}
}
}
Expand Down Expand Up @@ -1072,7 +1101,7 @@ unsafe impl BufMut for BytesMut {
if self.capacity() == self.len() {
self.reserve(64);
}
self.uninit_slice()
UninitSlice::from_slice(self.spare_capacity_mut())
}

// Specialize these methods so they can skip checking `remaining_mut`
Expand All @@ -1097,7 +1126,7 @@ unsafe impl BufMut for BytesMut {
fn put_bytes(&mut self, val: u8, cnt: usize) {
self.reserve(cnt);
unsafe {
let dst = self.uninit_slice();
let dst = self.spare_capacity_mut();
// Reserved above
debug_assert!(dst.len() >= cnt);

Expand Down