diff --git a/src/buf/buf_mut.rs b/src/buf/buf_mut.rs index ecf4720ba..d976ec419 100644 --- a/src/buf/buf_mut.rs +++ b/src/buf/buf_mut.rs @@ -1,10 +1,10 @@ #[cfg(feature = "std")] use super::Writer; -use core::{mem, cmp, ptr, usize}; +use core::{cmp, mem::{self, MaybeUninit}, ptr, usize}; #[cfg(feature = "std")] -use std::io::IoSliceMut; +use std::fmt; use alloc::{vec::Vec, boxed::Box}; @@ -72,13 +72,15 @@ pub trait BufMut { /// let mut buf = Vec::with_capacity(16); /// /// unsafe { - /// buf.bytes_mut()[0] = b'h'; - /// buf.bytes_mut()[1] = b'e'; + /// // MaybeUninit::as_mut_ptr + /// buf.bytes_mut()[0].as_mut_ptr().write(b'h'); + /// buf.bytes_mut()[1].as_mut_ptr().write(b'e'); /// /// buf.advance_mut(2); /// - /// buf.bytes_mut()[0] = b'l'; - /// buf.bytes_mut()[1..3].copy_from_slice(b"lo"); + /// buf.bytes_mut()[0].as_mut_ptr().write(b'l'); + /// buf.bytes_mut()[1].as_mut_ptr().write(b'l'); + /// buf.bytes_mut()[2].as_mut_ptr().write(b'o'); /// /// buf.advance_mut(3); /// } @@ -139,13 +141,15 @@ pub trait BufMut { /// let mut buf = Vec::with_capacity(16); /// /// unsafe { - /// buf.bytes_mut()[0] = b'h'; - /// buf.bytes_mut()[1] = b'e'; + /// // MaybeUninit::as_mut_ptr + /// buf.bytes_mut()[0].as_mut_ptr().write(b'h'); + /// buf.bytes_mut()[1].as_mut_ptr().write(b'e'); /// /// buf.advance_mut(2); /// - /// buf.bytes_mut()[0] = b'l'; - /// buf.bytes_mut()[1..3].copy_from_slice(b"lo"); + /// buf.bytes_mut()[0].as_mut_ptr().write(b'l'); + /// buf.bytes_mut()[1].as_mut_ptr().write(b'l'); + /// buf.bytes_mut()[2].as_mut_ptr().write(b'o'); /// /// buf.advance_mut(3); /// } @@ -161,7 +165,7 @@ pub trait BufMut { /// `bytes_mut` returning an empty slice implies that `remaining_mut` will /// return 0 and `remaining_mut` returning 0 implies that `bytes_mut` will /// return an empty slice. - unsafe fn bytes_mut(&mut self) -> &mut [u8]; + fn bytes_mut(&mut self) -> &mut [MaybeUninit]; /// Fills `dst` with potentially multiple mutable slices starting at `self`'s /// current position. @@ -192,13 +196,13 @@ pub trait BufMut { /// /// [`readv`]: http://man7.org/linux/man-pages/man2/readv.2.html #[cfg(feature = "std")] - unsafe fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize { + fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize { if dst.is_empty() { return 0; } if self.has_remaining_mut() { - dst[0] = IoSliceMut::new(self.bytes_mut()); + dst[0] = IoSliceMut::from(self.bytes_mut()); 1 } else { 0 @@ -238,7 +242,7 @@ pub trait BufMut { ptr::copy_nonoverlapping( s.as_ptr(), - d.as_mut_ptr(), + d.as_mut_ptr() as *mut u8, l); } @@ -280,7 +284,7 @@ pub trait BufMut { ptr::copy_nonoverlapping( src[off..].as_ptr(), - dst.as_mut_ptr(), + dst.as_mut_ptr() as *mut u8, cnt); off += cnt; @@ -931,12 +935,12 @@ impl BufMut for &mut T { (**self).remaining_mut() } - unsafe fn bytes_mut(&mut self) -> &mut [u8] { + fn bytes_mut(&mut self) -> &mut [MaybeUninit] { (**self).bytes_mut() } #[cfg(feature = "std")] - unsafe fn bytes_vectored_mut<'b>(&'b mut self, dst: &mut [IoSliceMut<'b>]) -> usize { + fn bytes_vectored_mut<'b>(&'b mut self, dst: &mut [IoSliceMut<'b>]) -> usize { (**self).bytes_vectored_mut(dst) } @@ -950,12 +954,12 @@ impl BufMut for Box { (**self).remaining_mut() } - unsafe fn bytes_mut(&mut self) -> &mut [u8] { + fn bytes_mut(&mut self) -> &mut [MaybeUninit] { (**self).bytes_mut() } #[cfg(feature = "std")] - unsafe fn bytes_vectored_mut<'b>(&'b mut self, dst: &mut [IoSliceMut<'b>]) -> usize { + fn bytes_vectored_mut<'b>(&'b mut self, dst: &mut [IoSliceMut<'b>]) -> usize { (**self).bytes_vectored_mut(dst) } @@ -971,8 +975,9 @@ impl BufMut for &mut [u8] { } #[inline] - unsafe fn bytes_mut(&mut self) -> &mut [u8] { - self + fn bytes_mut(&mut self) -> &mut [MaybeUninit] { + // MaybeUninit is repr(transparent), so safe to transmute + unsafe { mem::transmute(&mut **self) } } #[inline] @@ -1003,7 +1008,7 @@ impl BufMut for Vec { } #[inline] - unsafe fn bytes_mut(&mut self) -> &mut [u8] { + fn bytes_mut(&mut self) -> &mut [MaybeUninit] { use core::slice; if self.capacity() == self.len() { @@ -1013,11 +1018,54 @@ impl BufMut for Vec { let cap = self.capacity(); let len = self.len(); - let ptr = self.as_mut_ptr(); - &mut slice::from_raw_parts_mut(ptr, cap)[len..] + let ptr = self.as_mut_ptr() as *mut MaybeUninit; + unsafe { + &mut slice::from_raw_parts_mut(ptr, cap)[len..] + } } } // The existence of this function makes the compiler catch if the BufMut // trait is "object-safe" or not. fn _assert_trait_object(_b: &dyn BufMut) {} + +// ===== impl IoSliceMut ===== + +/// A buffer type used for `readv`. +/// +/// This is a wrapper around an `std::io::IoSliceMut`, but does not expose +/// the inner bytes in a safe API, as they may point at uninitialized memory. +/// +/// This is `repr(transparent)` of the `std::io::IoSliceMut`, so it is valid to +/// transmute them. However, as the memory might be uninitialized, care must be +/// taken to not *read* the internal bytes, only *write* to them. +#[repr(transparent)] +#[cfg(feature = "std")] +pub struct IoSliceMut<'a>(std::io::IoSliceMut<'a>); + +#[cfg(feature = "std")] +impl fmt::Debug for IoSliceMut<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("IoSliceMut") + .field("len", &self.0.len()) + .finish() + } +} + +#[cfg(feature = "std")] +impl<'a> From<&'a mut [u8]> for IoSliceMut<'a> { + fn from(buf: &'a mut [u8]) -> IoSliceMut<'a> { + IoSliceMut(std::io::IoSliceMut::new(buf)) + } +} + +#[cfg(feature = "std")] +impl<'a> From<&'a mut [MaybeUninit]> for IoSliceMut<'a> { + fn from(buf: &'a mut [MaybeUninit]) -> IoSliceMut<'a> { + IoSliceMut(std::io::IoSliceMut::new(unsafe { + // We don't look at the contents, and `std::io::IoSliceMut` + // doesn't either. + mem::transmute::<&'a mut [MaybeUninit], &'a mut [u8]>(buf) + })) + } +} diff --git a/src/buf/chain.rs b/src/buf/chain.rs index 7090ac49f..2805bf365 100644 --- a/src/buf/chain.rs +++ b/src/buf/chain.rs @@ -1,8 +1,12 @@ use crate::{Buf, BufMut}; use crate::buf::IntoIter; +use core::mem::MaybeUninit; + +#[cfg(feature = "std")] +use std::io::{IoSlice}; #[cfg(feature = "std")] -use std::io::{IoSlice, IoSliceMut}; +use crate::buf::IoSliceMut; /// A `Chain` sequences two buffers. /// @@ -196,7 +200,7 @@ impl BufMut for Chain self.a.remaining_mut() + self.b.remaining_mut() } - unsafe fn bytes_mut(&mut self) -> &mut [u8] { + fn bytes_mut(&mut self) -> &mut [MaybeUninit] { if self.a.has_remaining_mut() { self.a.bytes_mut() } else { @@ -223,7 +227,7 @@ impl BufMut for Chain } #[cfg(feature = "std")] - unsafe fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize { + fn bytes_vectored_mut<'a>(&'a mut self, dst: &mut [IoSliceMut<'a>]) -> usize { let mut n = self.a.bytes_vectored_mut(dst); n += self.b.bytes_vectored_mut(&mut dst[n..]); n diff --git a/src/buf/mod.rs b/src/buf/mod.rs index 15e032de0..1448abdea 100644 --- a/src/buf/mod.rs +++ b/src/buf/mod.rs @@ -31,6 +31,8 @@ mod writer; pub use self::buf_impl::Buf; pub use self::buf_mut::BufMut; +#[cfg(feature = "std")] +pub use self::buf_mut::IoSliceMut; pub use self::chain::Chain; pub use self::iter::IntoIter; #[cfg(feature = "std")] diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 6777929c6..8895602f7 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -926,19 +926,9 @@ impl BufMut for BytesMut { } #[inline] - unsafe fn bytes_mut(&mut self) -> &mut [u8] { - slice::from_raw_parts_mut(self.ptr.as_ptr().offset(self.len as isize), self.cap) - } - - #[inline] - fn put_slice(&mut self, src: &[u8]) { - assert!(self.remaining_mut() >= src.len()); - - let len = src.len(); - + fn bytes_mut(&mut self) -> &mut [mem::MaybeUninit] { unsafe { - self.bytes_mut()[..len].copy_from_slice(src); - self.advance_mut(len); + slice::from_raw_parts_mut(self.ptr.as_ptr().offset(self.len as isize) as *mut mem::MaybeUninit, self.cap) } } } @@ -1085,11 +1075,12 @@ impl Extend for BytesMut { let (lower, _) = iter.size_hint(); self.reserve(lower); + // TODO: optimize + // 1. If self.kind() == KIND_VEC, use Vec::extend + // 2. Make `reserve` inline-able for b in iter { - unsafe { - self.bytes_mut()[0] = b; - self.advance_mut(1); - } + self.reserve(1); + self.put_u8(b); } } } diff --git a/tests/test_buf_mut.rs b/tests/test_buf_mut.rs index 94654b7d2..0d372d1ec 100644 --- a/tests/test_buf_mut.rs +++ b/tests/test_buf_mut.rs @@ -1,9 +1,8 @@ #![deny(warnings, rust_2018_idioms)] -use bytes::{BufMut, BytesMut}; +use bytes::{buf::IoSliceMut, BufMut, BytesMut}; use std::usize; use std::fmt::Write; -use std::io::IoSliceMut; #[test] fn test_vec_as_mut_buf() { @@ -11,9 +10,7 @@ fn test_vec_as_mut_buf() { assert_eq!(buf.remaining_mut(), usize::MAX); - unsafe { - assert!(buf.bytes_mut().len() >= 64); - } + assert!(buf.bytes_mut().len() >= 64); buf.put(&b"zomg"[..]); @@ -72,20 +69,16 @@ fn test_clone() { fn test_bufs_vec_mut() { let b1: &mut [u8] = &mut []; let b2: &mut [u8] = &mut []; - let mut dst = [IoSliceMut::new(b1), IoSliceMut::new(b2)]; + let mut dst = [IoSliceMut::from(b1), IoSliceMut::from(b2)]; // with no capacity let mut buf = BytesMut::new(); assert_eq!(buf.capacity(), 0); - unsafe { - assert_eq!(0, buf.bytes_vectored_mut(&mut dst[..])); - } + assert_eq!(0, buf.bytes_vectored_mut(&mut dst[..])); // with capacity let mut buf = BytesMut::with_capacity(64); - unsafe { - assert_eq!(1, buf.bytes_vectored_mut(&mut dst[..])); - } + assert_eq!(1, buf.bytes_vectored_mut(&mut dst[..])); } #[test] diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 212189eea..9fa8019ee 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -464,6 +464,16 @@ fn extend_from_slice_mut() { } } +#[test] +fn extend_mut_without_size_hint() { + let mut bytes = BytesMut::with_capacity(0); + let mut long_iter = LONG.iter(); + + // Use iter::from_fn since it doesn't know a size_hint + bytes.extend(std::iter::from_fn(|| long_iter.next())); + assert_eq!(*bytes, LONG[..]); +} + #[test] fn from_static() { let mut a = Bytes::from_static(b"ab");