Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add Bytes::is_unique #643

Merged
merged 5 commits into from
Jan 19, 2024
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
50 changes: 50 additions & 0 deletions src/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,8 @@ pub(crate) struct Vtable {
///
/// takes `Bytes` to value
pub to_vec: unsafe fn(&AtomicPtr<()>, *const u8, usize) -> Vec<u8>,
/// fn(data)
pub is_unique: unsafe fn(&AtomicPtr<()>) -> bool,
/// fn(data, ptr, len)
pub drop: unsafe fn(&mut AtomicPtr<()>, *const u8, usize),
}
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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,
};

Expand All @@ -911,6 +936,10 @@ unsafe fn static_to_vec(_: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec<u8
slice.to_vec()
}

fn static_is_unique(_: &AtomicPtr<()>) -> bool {
false
}

unsafe fn static_drop(_: &mut AtomicPtr<()>, _: *const u8, _: usize) {
// nothing to drop for &'static [u8]
}
Expand All @@ -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,
};

Expand Down Expand Up @@ -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::<Shared>()).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())
Expand Down Expand Up @@ -1049,6 +1092,7 @@ const _: [(); 0 - mem::align_of::<Shared>() % 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,
};

Expand Down Expand Up @@ -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::<Shared>()).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());
Expand Down
1 change: 1 addition & 0 deletions src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down
33 changes: 33 additions & 0 deletions tests/test_bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<u8> = LONG.to_vec();
let b = Bytes::from(v);
assert!(b.is_unique());
}

#[test]
fn arc_is_unique() {
let v: Vec<u8> = 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<u8> = LONG.to_vec();
let b = Bytes::from(v);
let c = b.clone();
assert!(!c.is_unique());
drop(b);
assert!(c.is_unique());
}