diff --git a/src/bytes.rs b/src/bytes.rs index 9fed3d287..9eda9f463 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -112,6 +112,8 @@ pub(crate) struct Vtable { /// /// takes `Bytes` to value pub to_vec: unsafe fn(&AtomicPtr<()>, *const u8, usize) -> Vec, + /// fn(data) + pub is_unique: unsafe fn(&AtomicPtr<()>) -> bool, /// fn(data, ptr, len) pub drop: unsafe fn(&mut AtomicPtr<()>, *const u8, usize), } @@ -208,6 +210,28 @@ impl Bytes { self.len == 0 } + /// Returns true if this is the only reference to the data. + /// + /// Always returns false if the data is backed by a static slice. + /// + /// The result of this method may be invalidated immediately if another + /// thread clones this value while this is being called. Ensure you have + /// unique access to this value (`&mut Bytes`) first if you need to be + /// certain the result is valid (i.e. for safety reasons) + /// # Examples + /// + /// ``` + /// use bytes::Bytes; + /// + /// let a = Bytes::from(vec![1, 2, 3]); + /// assert!(a.is_unique()); + /// let b = a.clone(); + /// assert!(!a.is_unique()); + /// ``` + pub fn is_unique(&self) -> bool { + unsafe { (self.vtable.is_unique)(&self.data) } + } + /// Creates `Bytes` instance from slice, by copying it. pub fn copy_from_slice(data: &[u8]) -> Self { data.to_vec().into() @@ -898,6 +922,7 @@ impl fmt::Debug for Vtable { const STATIC_VTABLE: Vtable = Vtable { clone: static_clone, to_vec: static_to_vec, + is_unique: static_is_unique, drop: static_drop, }; @@ -911,6 +936,10 @@ unsafe fn static_to_vec(_: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec) -> bool { + false +} + unsafe fn static_drop(_: &mut AtomicPtr<()>, _: *const u8, _: usize) { // nothing to drop for &'static [u8] } @@ -920,12 +949,14 @@ 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, + is_unique: promotable_is_unique, drop: promotable_even_drop, }; static PROMOTABLE_ODD_VTABLE: Vtable = Vtable { clone: promotable_odd_clone, to_vec: promotable_odd_to_vec, + is_unique: promotable_is_unique, drop: promotable_odd_drop, }; @@ -1020,6 +1051,18 @@ unsafe fn promotable_odd_drop(data: &mut AtomicPtr<()>, ptr: *const u8, len: usi }); } +unsafe fn promotable_is_unique(data: &AtomicPtr<()>) -> bool { + let shared = data.load(Ordering::Acquire); + let kind = shared as usize & KIND_MASK; + + if kind == KIND_ARC { + let ref_cnt = (*shared.cast::()).ref_cnt.load(Ordering::Relaxed); + ref_cnt == 1 + } else { + true + } +} + unsafe fn free_boxed_slice(buf: *mut u8, offset: *const u8, len: usize) { let cap = (offset as usize - buf as usize) + len; dealloc(buf, Layout::from_size_align(cap, 1).unwrap()) @@ -1049,6 +1092,7 @@ const _: [(); 0 - mem::align_of::() % 2] = []; // Assert that the alignm static SHARED_VTABLE: Vtable = Vtable { clone: shared_clone, to_vec: shared_to_vec, + is_unique: shared_is_unique, drop: shared_drop, }; @@ -1094,6 +1138,12 @@ unsafe fn shared_to_vec(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec shared_to_vec_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); + ref_cnt == 1 +} + unsafe fn shared_drop(data: &mut AtomicPtr<()>, _ptr: *const u8, _len: usize) { data.with_mut(|shared| { release_shared(shared.cast()); diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 57fd33e4b..cb9687d9d 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1702,6 +1702,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, + is_unique: crate::bytes::shared_is_unique, drop: shared_v_drop, }; diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 5ec60a5b0..76adfdbf4 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -1208,3 +1208,36 @@ fn test_bytes_capacity_len() { } } } + +#[test] +fn static_is_unique() { + let b = Bytes::from_static(LONG); + assert!(!b.is_unique()); +} + +#[test] +fn vec_is_unique() { + let v: Vec = LONG.to_vec(); + let b = Bytes::from(v); + assert!(b.is_unique()); +} + +#[test] +fn arc_is_unique() { + let v: Vec = LONG.to_vec(); + let b = Bytes::from(v); + let c = b.clone(); + assert!(!b.is_unique()); + drop(c); + assert!(b.is_unique()); +} + +#[test] +fn shared_is_unique() { + let v: Vec = LONG.to_vec(); + let b = Bytes::from(v); + let c = b.clone(); + assert!(!c.is_unique()); + drop(b); + assert!(c.is_unique()); +}