Skip to content

Commit

Permalink
Add conversion from BytesMut to Vec<u8> (tokio-rs#543)
Browse files Browse the repository at this point in the history
  • Loading branch information
NobodyXu authored and lelongg committed Jan 9, 2023
1 parent 2575cb2 commit 30015c2
Show file tree
Hide file tree
Showing 2 changed files with 74 additions and 0 deletions.
37 changes: 37 additions & 0 deletions src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1540,6 +1540,43 @@ impl PartialEq<Bytes> for BytesMut {
}
}

impl From<BytesMut> for Vec<u8> {
fn from(mut bytes: BytesMut) -> Self {
let kind = bytes.kind();

let mut vec = if kind == KIND_VEC {
unsafe {
let (off, _) = bytes.get_vec_pos();
rebuild_vec(bytes.ptr.as_ptr(), bytes.len, bytes.cap, off)
}
} else if kind == KIND_ARC {
let shared = unsafe { &mut *(bytes.data as *mut Shared) };
if shared.is_unique() {
let vec = mem::replace(&mut shared.vec, Vec::new());

unsafe { release_shared(shared) };

vec
} else {
return bytes.deref().into();
}
} else {
return bytes.deref().into();
};

let len = bytes.len;

unsafe {
ptr::copy(bytes.ptr.as_ptr(), vec.as_mut_ptr(), len);
vec.set_len(len);
}

mem::forget(bytes);

vec
}
}

#[inline]
fn vptr(ptr: *mut u8) -> NonNull<u8> {
if cfg!(debug_assertions) {
Expand Down
37 changes: 37 additions & 0 deletions tests/test_bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1029,3 +1029,40 @@ fn box_slice_empty() {
let b = Bytes::from(empty);
assert!(b.is_empty());
}

#[test]
fn bytes_into_vec() {
// Test kind == KIND_VEC
let content = b"helloworld";

let mut bytes = BytesMut::new();
bytes.put_slice(content);

let vec: Vec<u8> = bytes.into();
assert_eq!(&vec, content);

// Test kind == KIND_ARC, shared.is_unique() == True
let mut bytes = BytesMut::new();
bytes.put_slice(b"abcdewe23");
bytes.put_slice(content);

// Overwrite the bytes to make sure only one reference to the underlying
// Vec exists.
bytes = bytes.split_off(9);

let vec: Vec<u8> = bytes.into();
assert_eq!(&vec, content);

// Test kind == KIND_ARC, shared.is_unique() == False
let prefix = b"abcdewe23";

let mut bytes = BytesMut::new();
bytes.put_slice(prefix);
bytes.put_slice(content);

let vec: Vec<u8> = bytes.split_off(prefix.len()).into();
assert_eq!(&vec, content);

let vec: Vec<u8> = bytes.into();
assert_eq!(&vec, prefix);
}

0 comments on commit 30015c2

Please sign in to comment.