diff --git a/src/bytes.rs b/src/bytes.rs index 9b8e279..0f1aaed 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -1,4 +1,4 @@ -use core::mem::{self, ManuallyDrop}; +use core::mem::{self, ManuallyDrop, MaybeUninit}; use core::ops::{Deref, RangeBounds}; use core::ptr::NonNull; use core::{cmp, fmt, hash, ptr, slice}; @@ -946,24 +946,26 @@ impl From<&'static str> for Bytes { impl From> for Bytes { fn from(vec: Vec) -> Bytes { + // Avoid an extra allocation if possible. + if vec.len() == vec.capacity() { + return Bytes::from(vec.into_boxed_slice()); + } + + let shared = Box::new(MaybeUninit::::uninit()); let mut vec = ManuallyDrop::new(vec); let ptr = vec.as_mut_ptr(); let len = vec.len(); let cap = vec.capacity(); - // Avoid an extra allocation if possible. - if len == cap { - let vec = ManuallyDrop::into_inner(vec); - return Bytes::from(vec.into_boxed_slice()); - } + let shared = Shared::init_to_raw( + shared, + Shared { + buf: ptr, + cap, + ref_cnt: AtomicUsize::new(1), + }, + ); - let shared = Box::new(Shared { - buf: ptr, - cap, - ref_cnt: AtomicUsize::new(1), - }); - - let shared = Box::into_raw(shared); // The pointer should be aligned, so this assert should // always succeed. debug_assert!( @@ -1327,6 +1329,15 @@ struct Shared { ref_cnt: AtomicUsize, } +impl Shared { + fn init_to_raw(b: Box>, v: Self) -> *mut Self { + let shared = Box::into_raw(b).cast::(); + // SAFETY: The Box has the right layout. + unsafe { shared.write(v) }; + shared + } +} + impl Drop for Shared { fn drop(&mut self) { unsafe { dealloc(self.buf, Layout::from_size_align(self.cap, 1).unwrap()) } @@ -1472,16 +1483,18 @@ unsafe fn shallow_clone_vec( // updated and since the buffer hasn't been promoted to an // `Arc`, those three fields still are the components of the // vector. - let shared = Box::new(Shared { - buf, - cap: offset.offset_from(buf) as usize + len, - // Initialize refcount to 2. One for this reference, and one - // for the new clone that will be returned from - // `shallow_clone`. - ref_cnt: AtomicUsize::new(2), - }); - - let shared = Box::into_raw(shared); + let shared = Box::new(MaybeUninit::::uninit()); + let shared = Shared::init_to_raw( + shared, + Shared { + buf, + cap: offset.offset_from(buf) as usize + len, + // Initialize refcount to 2. One for this reference, and one + // for the new clone that will be returned from + // `shallow_clone`. + ref_cnt: AtomicUsize::new(2), + }, + ); // The pointer should be aligned, so this assert should // always succeed. diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 689f4b4..b983bd1 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -79,6 +79,15 @@ struct Shared { ref_count: AtomicUsize, } +impl Shared { + fn init_to_raw(b: Box>, v: Self) -> *mut Self { + let shared = Box::into_raw(b).cast::(); + // SAFETY: The Box has the right layout. + unsafe { shared.write(v) }; + shared + } +} + // Assert that the alignment of `Shared` is divisible by 2. // This is a necessary invariant since we depend on allocating `Shared` a // shared object to implicitly carry the `KIND_ARC` flag in its pointer. @@ -1112,13 +1121,18 @@ impl BytesMut { // updated and since the buffer hasn't been promoted to an // `Arc`, those three fields still are the components of the // vector. - let shared = Box::new(Shared { - vec: rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off), - original_capacity_repr, - ref_count: AtomicUsize::new(ref_cnt), - }); - - let shared = Box::into_raw(shared); + // + // Explicitly allocate before invoking rebuild_vec() so that + // the vector is not dropped if Box::new() panics. + let shared = Box::new(MaybeUninit::::uninit()); + let shared = Shared::init_to_raw( + shared, + Shared { + vec: rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off), + original_capacity_repr, + ref_count: AtomicUsize::new(ref_cnt), + }, + ); // The pointer should be aligned, so this assert should // always succeed.