Add conversion from BytesMut to Vec<u8> (tokio-rs#543)

Author: NobodyXu

src/bytes_mut.rs

@@ -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) {

tests/test_bytes.rs

@@ -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);
+}