From 5dab74a6ced85924f69d1e1fcf2c52b34391e4a1 Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Sun, 14 Apr 2024 14:17:37 -0400 Subject: [PATCH 01/10] feat: Add zero-copy make_mut --- src/bytes.rs | 139 +++++++++++++++++++++++++++++++++++++++++++- src/bytes_mut.rs | 40 +++++++++++-- tests/test_bytes.rs | 88 ++++++++++++++++++++++++++++ 3 files changed, 261 insertions(+), 6 deletions(-) diff --git a/src/bytes.rs b/src/bytes.rs index 908cee9ad..9ea8d8596 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -15,7 +15,7 @@ use crate::buf::IntoIter; #[allow(unused)] use crate::loom::sync::atomic::AtomicMut; use crate::loom::sync::atomic::{AtomicPtr, AtomicUsize, Ordering}; -use crate::Buf; +use crate::{Buf, BytesMut}; /// A cheaply cloneable and sliceable chunk of contiguous memory. /// @@ -113,6 +113,7 @@ pub(crate) struct Vtable { /// /// takes `Bytes` to value pub to_vec: unsafe fn(&AtomicPtr<()>, *const u8, usize) -> Vec, + pub to_mut: unsafe fn(&AtomicPtr<()>, *const u8, usize) -> BytesMut, /// fn(data) pub is_unique: unsafe fn(&AtomicPtr<()>) -> bool, /// fn(data, ptr, len) @@ -507,6 +508,46 @@ impl Bytes { self.truncate(0); } + /// Try to convert self into `BytesMut`. + /// + /// If `self` is unique, this will succeed and return a `BytesMut` with + /// the contents of `self` without copying. If `self` is not unique, this + /// will fail and return self. + /// + /// # Examples + /// + /// ``` + /// use bytes::{Bytes, BytesMut}; + /// + /// let bytes = Bytes::from(b"hello".to_vec()); + /// assert_eq!(bytes.try_into_mut(), Ok(BytesMut::from(&b"hello"[..]))); + /// ``` + pub fn try_into_mut(self) -> Result { + if self.is_unique() { + Ok(self.make_mut()) + } else { + Err(self) + } + } + + /// Convert self into `BytesMut`. + /// + /// If `self` is unique, it will `BytesMut` with the contents of `self` without copying. + /// If `self` is not unique, it will make a copy of the data. + /// + /// # Examples + /// + /// ``` + /// use bytes::{Bytes, BytesMut}; + /// + /// let bytes = Bytes::from(b"hello".to_vec()); + /// assert_eq!(bytes.make_mut(), BytesMut::from(&b"hello"[..])); + /// ``` + pub fn make_mut(self) -> BytesMut { + let bytes = ManuallyDrop::new(self); + unsafe { (bytes.vtable.to_mut)(&bytes.data, bytes.ptr, bytes.len) } + } + #[inline] pub(crate) unsafe fn with_vtable( ptr: *const u8, @@ -917,6 +958,7 @@ impl fmt::Debug for Vtable { const STATIC_VTABLE: Vtable = Vtable { clone: static_clone, to_vec: static_to_vec, + to_mut: static_to_mut, is_unique: static_is_unique, drop: static_drop, }; @@ -931,6 +973,11 @@ unsafe fn static_to_vec(_: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec, ptr: *const u8, len: usize) -> BytesMut { + let slice = slice::from_raw_parts(ptr, len); + BytesMut::from(slice) +} + fn static_is_unique(_: &AtomicPtr<()>) -> bool { false } @@ -944,6 +991,7 @@ unsafe fn static_drop(_: &mut AtomicPtr<()>, _: *const u8, _: usize) { static PROMOTABLE_EVEN_VTABLE: Vtable = Vtable { clone: promotable_even_clone, to_vec: promotable_even_to_vec, + to_mut: promotable_even_to_mut, is_unique: promotable_is_unique, drop: promotable_even_drop, }; @@ -951,6 +999,7 @@ static PROMOTABLE_EVEN_VTABLE: Vtable = Vtable { static PROMOTABLE_ODD_VTABLE: Vtable = Vtable { clone: promotable_odd_clone, to_vec: promotable_odd_to_vec, + to_mut: promotable_odd_to_mut, is_unique: promotable_is_unique, drop: promotable_odd_drop, }; @@ -994,12 +1043,41 @@ unsafe fn promotable_to_vec( } } +unsafe fn promotable_to_mut( + data: &AtomicPtr<()>, + ptr: *const u8, + len: usize, + f: fn(*mut ()) -> *mut u8, +) -> BytesMut { + let shared = data.load(Ordering::Acquire); + let kind = shared as usize & KIND_MASK; + + if kind == KIND_ARC { + shared_to_mut_impl(shared.cast(), ptr, len) + } else { + debug_assert_eq!(kind, KIND_VEC); + + let buf = f(shared); + let off = offset_from(ptr, buf); + let cap = off + len; + let v = Vec::from_raw_parts(buf, cap, cap); + + BytesMut::from_vec_offset(v, off) + } +} + unsafe fn promotable_even_to_vec(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec { promotable_to_vec(data, ptr, len, |shared| { ptr_map(shared.cast(), |addr| addr & !KIND_MASK) }) } +unsafe fn promotable_even_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> BytesMut { + promotable_to_mut(data, ptr, len, |shared| { + ptr_map(shared.cast(), |addr| addr & !KIND_MASK) + }) +} + unsafe fn promotable_even_drop(data: &mut AtomicPtr<()>, ptr: *const u8, len: usize) { data.with_mut(|shared| { let shared = *shared; @@ -1031,6 +1109,10 @@ unsafe fn promotable_odd_to_vec(data: &AtomicPtr<()>, ptr: *const u8, len: usize promotable_to_vec(data, ptr, len, |shared| shared.cast()) } +unsafe fn promotable_odd_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> BytesMut { + promotable_to_mut(data, ptr, len, |shared| shared.cast()) +} + unsafe fn promotable_odd_drop(data: &mut AtomicPtr<()>, ptr: *const u8, len: usize) { data.with_mut(|shared| { let shared = *shared; @@ -1087,6 +1169,7 @@ const _: [(); 0 - mem::align_of::() % 2] = []; // Assert that the alignm static SHARED_VTABLE: Vtable = Vtable { clone: shared_clone, to_vec: shared_to_vec, + to_mut: shared_to_mut, is_unique: shared_is_unique, drop: shared_drop, }; @@ -1133,6 +1216,43 @@ unsafe fn shared_to_vec(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec shared_to_vec_impl(data.load(Ordering::Relaxed).cast(), ptr, len) } +unsafe fn shared_to_mut_impl(shared: *mut Shared, ptr: *const u8, len: usize) -> BytesMut { + // The goal is to check if the current handle is the only handle + // that currently has access to the buffer. This is done by + // checking if the `ref_cnt` is currently 1. + // + // The `Acquire` ordering synchronizes with the `Release` as + // part of the `fetch_sub` in `release_shared`. The `fetch_sub` + // operation guarantees that any mutations done in other threads + // are ordered before the `ref_cnt` is decremented. As such, + // this `Acquire` will guarantee that those mutations are + // visible to the current thread. + // + // Otherwise, we take the other branch, copy the data and call `release_shared``. + if (*shared).ref_cnt.load(Ordering::Acquire) == 1 { + // Deallocate the `Shared` instance without running its destructor. + let shared = *Box::from_raw(shared); + let shared = ManuallyDrop::new(shared); + let buf = shared.buf; + let cap = shared.cap; + + // Rebuild Vec + let off = offset_from(ptr, buf); + let len = len + off; + let v = Vec::from_raw_parts(buf, len, cap); + + BytesMut::from_vec_offset(v, off) + } else { + let v = slice::from_raw_parts(ptr, len).to_vec(); + release_shared(shared); + BytesMut::from_vec(v) + } +} + +unsafe fn shared_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> BytesMut { + shared_to_mut_impl(data.load(Ordering::Relaxed).cast(), ptr, len) +} + pub(crate) unsafe fn shared_is_unique(data: &AtomicPtr<()>) -> bool { let shared = data.load(Ordering::Acquire); let ref_cnt = (*shared.cast::()).ref_cnt.load(Ordering::Relaxed); @@ -1291,6 +1411,23 @@ where new_addr as *mut u8 } +/// Precondition: dst >= original +/// +/// The following line is equivalent to: +/// +/// ```rust,ignore +/// self.ptr.as_ptr().offset_from(ptr) as usize; +/// ``` +/// +/// But due to min rust is 1.39 and it is only stabilized +/// in 1.47, we cannot use it. +#[inline] +fn offset_from(dst: *const u8, original: *const u8) -> usize { + debug_assert!(dst >= original); + + dst as usize - original as usize +} + // compile-fails /// ```compile_fail diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index b01bb1adc..5f7e5976a 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -837,13 +837,20 @@ impl BytesMut { // suddenly a lot more expensive. #[inline] pub(crate) fn from_vec(vec: Vec) -> BytesMut { + unsafe { BytesMut::from_vec_offset(vec, 0) } + } + + #[inline] + pub(crate) unsafe fn from_vec_offset(vec: Vec, off: usize) -> BytesMut { let mut vec = ManuallyDrop::new(vec); - let ptr = vptr(vec.as_mut_ptr()); - let len = vec.len(); - let cap = vec.capacity(); + let ptr = vptr(vec.as_mut_ptr().add(off)); + let len = vec.len().checked_sub(off).unwrap_or(0); + let cap = vec.capacity() - off; - let original_capacity_repr = original_capacity_to_repr(cap); - let data = (original_capacity_repr << ORIGINAL_CAPACITY_OFFSET) | KIND_VEC; + let original_capacity_repr = original_capacity_to_repr(vec.capacity()); + let data = (original_capacity_repr << ORIGINAL_CAPACITY_OFFSET) + | (off << VEC_POS_OFFSET) + | KIND_VEC; BytesMut { ptr, @@ -1713,6 +1720,7 @@ unsafe fn rebuild_vec(ptr: *mut u8, mut len: usize, mut cap: usize, off: usize) static SHARED_VTABLE: Vtable = Vtable { clone: shared_v_clone, to_vec: shared_v_to_vec, + to_mut: shared_v_to_mut, is_unique: crate::bytes::shared_is_unique, drop: shared_v_drop, }; @@ -1747,6 +1755,28 @@ unsafe fn shared_v_to_vec(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> V } } +unsafe fn shared_v_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> BytesMut { + let shared: *mut Shared = data.load(Ordering::Relaxed).cast(); + + if (*shared).is_unique() { + let shared = &mut *shared; + + let ptr = vptr(ptr as *mut u8); + let cap = len; // It will try to reclaim the buffer on first insert + + BytesMut { + ptr, + len, + cap, + data: shared, + } + } else { + let v = slice::from_raw_parts(ptr, len).to_vec(); + release_shared(shared); + BytesMut::from_vec(v) + } +} + unsafe fn shared_v_drop(data: &mut AtomicPtr<()>, _ptr: *const u8, _len: usize) { data.with_mut(|shared| { release_shared(*shared as *mut Shared); diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 84c3d5a43..2f70b95ee 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -1172,3 +1172,91 @@ fn shared_is_unique() { drop(b); assert!(c.is_unique()); } + +#[test] +fn test_bytes_make_mut() { + // Test STATIC_VTABLE.to_mut + let bs = b"1b23exfcz3r"; + let bs_long = b"1b23exfcz3r1b23exfcz3r"; + let bytes_mut = Bytes::from_static(bs).make_mut(); + assert_eq!(bytes_mut, bs[..]); + + // Test bytes_mut.SHARED_VTABLE.to_mut impl + eprintln!("1"); + let mut bytes_mut: BytesMut = bs[..].into(); + bytes_mut = bytes_mut.freeze().make_mut(); + assert_eq!(bytes_mut, bs[..]); + bytes_mut.extend_from_slice(&bs[..]); + assert_eq!(bytes_mut, bs_long[..]); + + // Set kind to KIND_ARC so that after freeze, Bytes will use bytes_mut.SHARED_VTABLE + let mut bytes_mut: BytesMut = bs[..].into(); + eprintln!("2"); + drop(bytes_mut.split_off(bs.len())); + + eprintln!("3"); + let b1 = bytes_mut.freeze(); + eprintln!("4"); + let b2 = b1.clone(); + + eprintln!("{:#?}", (&*b1).as_ptr()); + + // shared.is_unique() = False + eprintln!("5"); + let mut b1m = b1.make_mut(); + assert_eq!(b1m, bs[..]); + b1m[0] = b'9'; + + // shared.is_unique() = True + eprintln!("6"); + let b2m = b2.make_mut(); + assert_eq!(b2m, bs[..]); + + // Test bytes_mut.SHARED_VTABLE.to_mut impl where offset != 0 + let mut bytes_mut1: BytesMut = bs[..].into(); + let bytes_mut2 = bytes_mut1.split_off(9); + + let b1 = bytes_mut1.freeze(); + let b2 = bytes_mut2.freeze(); + + let b1m = b1.make_mut(); + let b2m = b2.make_mut(); + + assert_eq!(b2m, bs[9..]); + assert_eq!(b1m, bs[..9]); +} + +#[test] +fn test_bytes_make_mut_promotable_even() { + let vec = vec![33u8; 1024]; + + // Test cases where kind == KIND_VEC + let b1 = Bytes::from(vec.clone()); + let b1m = b1.make_mut(); + assert_eq!(b1m, vec); + + // Test cases where kind == KIND_ARC, ref_cnt == 1 + let b1 = Bytes::from(vec.clone()); + drop(b1.clone()); + let b1m = b1.make_mut(); + assert_eq!(b1m, vec); + + // Test cases where kind == KIND_ARC, ref_cnt == 2 + let b1 = Bytes::from(vec.clone()); + let b2 = b1.clone(); + let b1m = b1.make_mut(); + assert_eq!(b1m, vec); + + // Test cases where vtable = SHARED_VTABLE, kind == KIND_ARC, ref_cnt == 1 + let b2m = b2.make_mut(); + assert_eq!(b2m, vec); + + // Test cases where offset != 0 + let mut b1 = Bytes::from(vec.clone()); + let b2 = b1.split_off(20); + let b1m = b1.make_mut(); + let b2m = b2.make_mut(); + + assert_eq!(b2m, vec[20..]); + assert_eq!(b1m, vec[..20]); +} From d75c2a2d7a3ba1eb6b261b0813a86cb539869f37 Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Tue, 30 Apr 2024 11:38:35 -0400 Subject: [PATCH 02/10] use checked sub to avoid overflow in cap --- src/bytes_mut.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 5f7e5976a..f8f892a90 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -845,7 +845,7 @@ impl BytesMut { let mut vec = ManuallyDrop::new(vec); let ptr = vptr(vec.as_mut_ptr().add(off)); let len = vec.len().checked_sub(off).unwrap_or(0); - let cap = vec.capacity() - off; + let cap = vec.capacity().checked_sub(off).unwrap_or(0); let original_capacity_repr = original_capacity_to_repr(vec.capacity()); let data = (original_capacity_repr << ORIGINAL_CAPACITY_OFFSET) From f6a633a79c042804ad5fb6c243b9cfe1ba2448b0 Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Tue, 30 Apr 2024 11:38:51 -0400 Subject: [PATCH 03/10] Add comment for KIND_VEC memory safety --- src/bytes.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/bytes.rs b/src/bytes.rs index 9ea8d8596..3de708927 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -1055,6 +1055,9 @@ unsafe fn promotable_to_mut( if kind == KIND_ARC { shared_to_mut_impl(shared.cast(), ptr, len) } else { + // The kind KIND_VEC always represents the full view of the underlying buffer, + // if it truncated it is first promoted to KIND_ARC. Thus we can reconstruct + // a Vec from it without leaking memory. debug_assert_eq!(kind, KIND_VEC); let buf = f(shared); From eee880cb7588cd2af68e0c883b240265bb622abe Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Tue, 30 Apr 2024 11:41:56 -0400 Subject: [PATCH 04/10] Add odd alloc make_mut test --- tests/test_bytes_odd_alloc.rs | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/test_bytes_odd_alloc.rs b/tests/test_bytes_odd_alloc.rs index 27ed87736..ae3df4cec 100644 --- a/tests/test_bytes_odd_alloc.rs +++ b/tests/test_bytes_odd_alloc.rs @@ -95,3 +95,38 @@ fn test_bytes_into_vec() { assert_eq!(Vec::from(b2), vec[20..]); assert_eq!(Vec::from(b1), vec[..20]); } + +#[test] +fn test_bytes_make_mut() { + let vec = vec![33u8; 1024]; + + // Test cases where kind == KIND_VEC + let b1 = Bytes::from(vec.clone()); + let b1m = b1.make_mut(); + assert_eq!(b1m, vec); + + // Test cases where kind == KIND_ARC, ref_cnt == 1 + let b1 = Bytes::from(vec.clone()); + drop(b1.clone()); + let b1m = b1.make_mut(); + assert_eq!(b1m, vec); + + // Test cases where kind == KIND_ARC, ref_cnt == 2 + let b1 = Bytes::from(vec.clone()); + let b2 = b1.clone(); + let b1m = b1.make_mut(); + assert_eq!(b1m, vec); + + // Test cases where vtable = SHARED_VTABLE, kind == KIND_ARC, ref_cnt == 1 + let b2m = b2.make_mut(); + assert_eq!(b2m, vec); + + // Test cases where offset != 0 + let mut b1 = Bytes::from(vec.clone()); + let b2 = b1.split_off(20); + let b1m = b1.make_mut(); + let b2m = b2.make_mut(); + + assert_eq!(b2m, vec[20..]); + assert_eq!(b1m, vec[..20]); +} From c5b07a55897c93041ae5f77b8656ef0ac6f7b5d9 Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Wed, 1 May 2024 14:17:55 -0400 Subject: [PATCH 05/10] Split tests --- tests/test_bytes.rs | 59 ++++++++++++++++++++++++----------- tests/test_bytes_odd_alloc.rs | 27 ++++++++++++---- 2 files changed, 62 insertions(+), 24 deletions(-) diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 2f70b95ee..2f283af2f 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -1174,43 +1174,51 @@ fn shared_is_unique() { } #[test] -fn test_bytes_make_mut() { - // Test STATIC_VTABLE.to_mut +fn test_bytes_make_mut_static() { let bs = b"1b23exfcz3r"; - let bs_long = b"1b23exfcz3r1b23exfcz3r"; + + // Test STATIC_VTABLE.to_mut let bytes_mut = Bytes::from_static(bs).make_mut(); assert_eq!(bytes_mut, bs[..]); +} - // Test bytes_mut.SHARED_VTABLE.to_mut impl - eprintln!("1"); +#[test] +fn test_bytes_make_mut_bytes_mut_vec() { + let bs = b"1b23exfcz3r"; + let bs_long = b"1b23exfcz3r1b23exfcz3r"; + + // Test case where kind == KIND_VEC let mut bytes_mut: BytesMut = bs[..].into(); bytes_mut = bytes_mut.freeze().make_mut(); assert_eq!(bytes_mut, bs[..]); bytes_mut.extend_from_slice(&bs[..]); assert_eq!(bytes_mut, bs_long[..]); +} + +#[test] +fn test_bytes_make_mut_bytes_mut_shared() { + let bs = b"1b23exfcz3r"; // Set kind to KIND_ARC so that after freeze, Bytes will use bytes_mut.SHARED_VTABLE let mut bytes_mut: BytesMut = bs[..].into(); - eprintln!("2"); drop(bytes_mut.split_off(bs.len())); - eprintln!("3"); let b1 = bytes_mut.freeze(); - eprintln!("4"); let b2 = b1.clone(); - eprintln!("{:#?}", (&*b1).as_ptr()); - // shared.is_unique() = False - eprintln!("5"); let mut b1m = b1.make_mut(); assert_eq!(b1m, bs[..]); b1m[0] = b'9'; // shared.is_unique() = True - eprintln!("6"); let b2m = b2.make_mut(); assert_eq!(b2m, bs[..]); +} + +#[test] +fn test_bytes_make_mut_bytes_mut_offset() { + let bs = b"1b23exfcz3r"; // Test bytes_mut.SHARED_VTABLE.to_mut impl where offset != 0 let mut bytes_mut1: BytesMut = bs[..].into(); @@ -1227,31 +1235,46 @@ fn test_bytes_make_mut() { } #[test] -fn test_bytes_make_mut_promotable_even() { +fn test_bytes_make_mut_promotable_even_vec() { let vec = vec![33u8; 1024]; - // Test cases where kind == KIND_VEC + // Test case where kind == KIND_VEC let b1 = Bytes::from(vec.clone()); let b1m = b1.make_mut(); assert_eq!(b1m, vec); +} - // Test cases where kind == KIND_ARC, ref_cnt == 1 +#[test] +fn test_bytes_make_mut_promotable_even_arc_1() { + let vec = vec![33u8; 1024]; + + // Test case where kind == KIND_ARC, ref_cnt == 1 let b1 = Bytes::from(vec.clone()); drop(b1.clone()); let b1m = b1.make_mut(); assert_eq!(b1m, vec); +} - // Test cases where kind == KIND_ARC, ref_cnt == 2 +#[test] +fn test_bytes_make_mut_promotable_even_arc_2() { + let vec = vec![33u8; 1024]; + + // Test case where kind == KIND_ARC, ref_cnt == 2 let b1 = Bytes::from(vec.clone()); let b2 = b1.clone(); let b1m = b1.make_mut(); assert_eq!(b1m, vec); - // Test cases where vtable = SHARED_VTABLE, kind == KIND_ARC, ref_cnt == 1 + // Test case where vtable = SHARED_VTABLE, kind == KIND_ARC, ref_cnt == 1 let b2m = b2.make_mut(); assert_eq!(b2m, vec); +} - // Test cases where offset != 0 +#[test] +fn test_bytes_make_mut_promotable_even_arc_offset() { + let vec = vec![33u8; 1024]; + + // Test case where offset != 0 let mut b1 = Bytes::from(vec.clone()); let b2 = b1.split_off(20); let b1m = b1.make_mut(); diff --git a/tests/test_bytes_odd_alloc.rs b/tests/test_bytes_odd_alloc.rs index ae3df4cec..8008a0e47 100644 --- a/tests/test_bytes_odd_alloc.rs +++ b/tests/test_bytes_odd_alloc.rs @@ -97,31 +97,46 @@ fn test_bytes_into_vec() { } #[test] -fn test_bytes_make_mut() { +fn test_bytes_make_mut_vec() { let vec = vec![33u8; 1024]; - // Test cases where kind == KIND_VEC + // Test case where kind == KIND_VEC let b1 = Bytes::from(vec.clone()); let b1m = b1.make_mut(); assert_eq!(b1m, vec); +} - // Test cases where kind == KIND_ARC, ref_cnt == 1 +#[test] +fn test_bytes_make_mut_arc_1() { + let vec = vec![33u8; 1024]; + + // Test case where kind == KIND_ARC, ref_cnt == 1 let b1 = Bytes::from(vec.clone()); drop(b1.clone()); let b1m = b1.make_mut(); assert_eq!(b1m, vec); +} - // Test cases where kind == KIND_ARC, ref_cnt == 2 +#[test] +fn test_bytes_make_mut_arc_2() { + let vec = vec![33u8; 1024]; + + // Test case where kind == KIND_ARC, ref_cnt == 2 let b1 = Bytes::from(vec.clone()); let b2 = b1.clone(); let b1m = b1.make_mut(); assert_eq!(b1m, vec); - // Test cases where vtable = SHARED_VTABLE, kind == KIND_ARC, ref_cnt == 1 + // Test case where vtable = SHARED_VTABLE, kind == KIND_ARC, ref_cnt == 1 let b2m = b2.make_mut(); assert_eq!(b2m, vec); +} - // Test cases where offset != 0 +#[test] +fn test_bytes_make_mut_arc_offset() { + let vec = vec![33u8; 1024]; + + // Test case where offset != 0 let mut b1 = Bytes::from(vec.clone()); let b2 = b1.split_off(20); let b1m = b1.make_mut(); From bb6a8db4c24eddaf0326e86161e87f92a33e3c30 Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Wed, 1 May 2024 14:18:01 -0400 Subject: [PATCH 06/10] Improve comments --- src/bytes.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/bytes.rs b/src/bytes.rs index 3de708927..de3c512b7 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -1055,9 +1055,10 @@ unsafe fn promotable_to_mut( if kind == KIND_ARC { shared_to_mut_impl(shared.cast(), ptr, len) } else { - // The kind KIND_VEC always represents the full view of the underlying buffer, - // if it truncated it is first promoted to KIND_ARC. Thus we can reconstruct - // a Vec from it without leaking memory. + // KIND_VEC is a view of an underlying buffer at a certain offset. + // The ptr + len always represents the end of that buffer. + // Before truncating it, it is first promoted to KIND_ARC. + // Thus, we can safely reconstruct a Vec from it without leaking memory. debug_assert_eq!(kind, KIND_VEC); let buf = f(shared); @@ -1231,7 +1232,7 @@ unsafe fn shared_to_mut_impl(shared: *mut Shared, ptr: *const u8, len: usize) -> // this `Acquire` will guarantee that those mutations are // visible to the current thread. // - // Otherwise, we take the other branch, copy the data and call `release_shared``. + // Otherwise, we take the other branch, copy the data and call `release_shared`. if (*shared).ref_cnt.load(Ordering::Acquire) == 1 { // Deallocate the `Shared` instance without running its destructor. let shared = *Box::from_raw(shared); @@ -1246,6 +1247,7 @@ unsafe fn shared_to_mut_impl(shared: *mut Shared, ptr: *const u8, len: usize) -> BytesMut::from_vec_offset(v, off) } else { + // Copy the data from Shared in a new Vec, then release it let v = slice::from_raw_parts(ptr, len).to_vec(); release_shared(shared); BytesMut::from_vec(v) From 3a0e2a4b0b4d81bcd2b650776da7fb8433cf9b09 Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Wed, 1 May 2024 14:33:44 -0400 Subject: [PATCH 07/10] Improve BytesMut SHARED capacity computing --- src/bytes_mut.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index f8f892a90..3f90ac377 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1761,8 +1761,15 @@ unsafe fn shared_v_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> B if (*shared).is_unique() { let shared = &mut *shared; + // The capacity of the buffer is always the original capacity + // of the buffer minus the offset from the start of that buffer. + let v = &mut shared.vec; + let v_capacity = v.capacity(); + let v_ptr = v.as_mut_ptr(); + let offset = offset_from(ptr as *mut u8, v_ptr); + let cap = v_capacity.checked_sub(offset).unwrap_or(len); + let ptr = vptr(ptr as *mut u8); - let cap = len; // It will try to reclaim the buffer on first insert BytesMut { ptr, From 94acd5e170e338602ed903dfcb50fc5badb426c9 Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Wed, 1 May 2024 14:39:06 -0400 Subject: [PATCH 08/10] Improve comment on capacity calculation --- src/bytes_mut.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 3f90ac377..329e68d58 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1761,8 +1761,8 @@ unsafe fn shared_v_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> B if (*shared).is_unique() { let shared = &mut *shared; - // The capacity of the buffer is always the original capacity - // of the buffer minus the offset from the start of that buffer. + // The capacity is always the original capacity of the buffer + // minus the offset from the start of the buffer let v = &mut shared.vec; let v_capacity = v.capacity(); let v_ptr = v.as_mut_ptr(); From ba06c9a77ee866a1062513d1dc3186d593f6bf0f Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Sat, 4 May 2024 17:21:36 -0400 Subject: [PATCH 09/10] Remove from_vec_offset --- src/bytes.rs | 11 +++++++---- src/bytes_mut.rs | 19 ++++++------------- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/src/bytes.rs b/src/bytes.rs index de3c512b7..065252787 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -1066,7 +1066,9 @@ unsafe fn promotable_to_mut( let cap = off + len; let v = Vec::from_raw_parts(buf, cap, cap); - BytesMut::from_vec_offset(v, off) + let mut b = BytesMut::from_vec(v); + b.advance_unchecked(off); + b } } @@ -1242,10 +1244,11 @@ unsafe fn shared_to_mut_impl(shared: *mut Shared, ptr: *const u8, len: usize) -> // Rebuild Vec let off = offset_from(ptr, buf); - let len = len + off; - let v = Vec::from_raw_parts(buf, len, cap); + let v = Vec::from_raw_parts(buf, len + off, cap); - BytesMut::from_vec_offset(v, off) + let mut b = BytesMut::from_vec(v); + b.advance_unchecked(off); + b } else { // Copy the data from Shared in a new Vec, then release it let v = slice::from_raw_parts(ptr, len).to_vec(); diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 329e68d58..686f4dd13 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -837,20 +837,13 @@ impl BytesMut { // suddenly a lot more expensive. #[inline] pub(crate) fn from_vec(vec: Vec) -> BytesMut { - unsafe { BytesMut::from_vec_offset(vec, 0) } - } - - #[inline] - pub(crate) unsafe fn from_vec_offset(vec: Vec, off: usize) -> BytesMut { let mut vec = ManuallyDrop::new(vec); - let ptr = vptr(vec.as_mut_ptr().add(off)); - let len = vec.len().checked_sub(off).unwrap_or(0); - let cap = vec.capacity().checked_sub(off).unwrap_or(0); + let ptr = vptr(vec.as_mut_ptr()); + let len = vec.len(); + let cap = vec.capacity(); - let original_capacity_repr = original_capacity_to_repr(vec.capacity()); - let data = (original_capacity_repr << ORIGINAL_CAPACITY_OFFSET) - | (off << VEC_POS_OFFSET) - | KIND_VEC; + let original_capacity_repr = original_capacity_to_repr(cap); + let data = (original_capacity_repr << ORIGINAL_CAPACITY_OFFSET) | KIND_VEC; BytesMut { ptr, @@ -875,7 +868,7 @@ impl BytesMut { /// # SAFETY /// /// The caller must ensure that `count` <= `self.cap`. - unsafe fn advance_unchecked(&mut self, count: usize) { + pub(crate) unsafe fn advance_unchecked(&mut self, count: usize) { // Setting the start to 0 is a no-op, so return early if this is the // case. if count == 0 { From 17fd79d26a159afc54fc38ed542fc07c8bcc21b6 Mon Sep 17 00:00:00 2001 From: Emile Fugulin Date: Sun, 5 May 2024 10:37:15 -0400 Subject: [PATCH 10/10] Improve comment for make_mut and remove unecessary checked_sub --- src/bytes.rs | 13 ++++++++----- src/bytes_mut.rs | 2 +- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/bytes.rs b/src/bytes.rs index 065252787..b4359b08d 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -510,9 +510,10 @@ impl Bytes { /// Try to convert self into `BytesMut`. /// - /// If `self` is unique, this will succeed and return a `BytesMut` with - /// the contents of `self` without copying. If `self` is not unique, this - /// will fail and return self. + /// If `self` is unique for the entire original buffer, this will succeed + /// and return a `BytesMut` with the contents of `self` without copying. + /// If `self` is not unique for the entire original buffer, this will fail + /// and return self. /// /// # Examples /// @@ -532,8 +533,10 @@ impl Bytes { /// Convert self into `BytesMut`. /// - /// If `self` is unique, it will `BytesMut` with the contents of `self` without copying. - /// If `self` is not unique, it will make a copy of the data. + /// If `self` is unique for the entire original buffer, this will return a + /// `BytesMut` with the contents of `self` without copying. + /// If `self` is not unique for the entire original buffer, this will make + /// a copy of `self` subset of the original buffer in a new `BytesMut`. /// /// # Examples /// diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 686f4dd13..569f8be63 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1760,7 +1760,7 @@ unsafe fn shared_v_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> B let v_capacity = v.capacity(); let v_ptr = v.as_mut_ptr(); let offset = offset_from(ptr as *mut u8, v_ptr); - let cap = v_capacity.checked_sub(offset).unwrap_or(len); + let cap = v_capacity - offset; let ptr = vptr(ptr as *mut u8);