Implement Bytes::unsplit and Bytes::try_unsplit

src/bytes.rs

@@ -172,7 +172,7 @@ impl Bytes {
         Bytes {
             ptr: bytes.as_ptr(),
             len: bytes.len(),
-            data: AtomicPtr::new(ptr::null_mut()),
+            data: AtomicPtr::new(bytes.as_ptr() as *mut ()),
             vtable: &STATIC_VTABLE,
         }
     }
@@ -183,7 +183,7 @@ impl Bytes {
         Bytes {
             ptr: bytes.as_ptr(),
             len: bytes.len(),
-            data: AtomicPtr::new(ptr::null_mut()),
+            data: AtomicPtr::new(bytes.as_ptr() as *mut ()),
             vtable: &STATIC_VTABLE,
         }
     }
@@ -199,7 +199,7 @@ impl Bytes {
         Bytes {
             ptr,
             len: 0,
-            data: AtomicPtr::new(ptr::null_mut()),
+            data: AtomicPtr::new(ptr as *mut ()),
             vtable: &STATIC_VTABLE,
         }
     }
@@ -547,6 +547,111 @@ impl Bytes {
         ret
     }
 
+    /// Absorbs a `Bytes` that was previously split off if they are contiguous;
+    /// otherwise, appends its bytes to this `Bytes` and return a new one.
+    ///
+    /// If the two `Bytes` objects were previously contiguous (i.e., if `other`
+    /// was created by calling `split_off` on this `Bytes`), then this is an
+    /// `O(1)` operation that just decreases a reference count and sets a few
+    /// indices.
+    ///
+    /// Otherwise, this method tries to convert `self` into a `BytesMut` and
+    /// extend it with `other`. If that also fails, this method allocates a new
+    /// `BytesMut` and copies both `Bytes` objects into it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::Bytes;
+    ///
+    /// let mut buf = Bytes::from_static(b"aaabbbcccddd");
+    ///
+    /// let split = buf.split_off(6);
+    /// assert_eq!(b"aaabbb", &buf[..]);
+    /// assert_eq!(b"cccddd", &split[..]);
+    ///
+    /// let buf = buf.unsplit(split);
+    /// assert_eq!(b"aaabbbcccddd", &buf[..]);
+    /// ```
+    pub fn unsplit(mut self, other: Bytes) -> Self {
+        if other.is_empty() {
+            return self;
+        } else if self.is_empty() {
+            return other;
+        }
+        // try to unsplit first
+        if let Err(other) = self.try_unsplit(other) {
+            // if failed, try to convert into a mut
+            let mut_self = match self.try_into_mut() {
+                Ok(mut mut_self) => {
+                    // if success, extend from other slice
+                    mut_self.extend_from_slice(&other);
+                    mut_self
+                }
+                Err(this) => {
+                    // if not, create a buffer with enough capacity
+                    let mut mut_self = BytesMut::with_capacity(this.len() + other.len());
+                    mut_self.extend_from_slice(&this);
+                    mut_self.extend_from_slice(&other);
+                    mut_self
+                }
+            };
+            // then freeze the mut into a bytes
+            mut_self.freeze()
+        } else {
+            self
+        }
+    }
+
+    /// Absorbs a `Bytes` that was previously split off.
+    ///
+    /// If the two `Bytes` objects were previously contiguous, i.e., if `other`
+    /// was created by calling `split_off` on this `Bytes`, then this is an
+    /// `O(1)` operation that just decreases a reference count and sets a few
+    /// indices. Otherwise, this method returns an error containing the
+    /// original `other`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::Bytes;
+    ///
+    /// let mut buf = Bytes::from_static(b"aaabbbcccddd");
+    ///
+    /// let mut split_1 = buf.split_off(3);
+    /// let split_2 = split_1.split_off(3);
+    /// assert_eq!(b"aaa", &buf[..]);
+    /// assert_eq!(b"bbb", &split_1[..]);
+    /// assert_eq!(b"cccddd", &split_2[..]);
+    ///
+    /// let split_2 = buf.try_unsplit(split_2).unwrap_err();
+    ///
+    /// buf.try_unsplit(split_1).unwrap();
+    /// buf.try_unsplit(split_2).unwrap();
+    /// assert_eq!(b"aaabbbcccddd", &buf[..]);
+    /// ```
+    pub fn try_unsplit(&mut self, other: Bytes) -> Result<(), Bytes> {
+        if other.len() == 0 {
+            return Ok(());
+        }
+
+        // Check if this block is right next to the other block
+        let ptr = unsafe { self.ptr.add(self.len) };
+        // ... and from the same backing data
+        if ptr == other.ptr && self.is_same_data(&other) {
+            // ... then combine two blocks into one
+            self.len += other.len;
+            Ok(())
+        } else {
+            Err(other)
+        }
+    }
+
+    #[inline]
+    fn is_same_data(&self, other: &Bytes) -> bool {
+        self.data.load(Ordering::Relaxed) == other.data.load(Ordering::Relaxed)
+    }
+
     /// Shortens the buffer, keeping the first `len` bytes and dropping the
     /// rest.
     ///
@@ -1072,9 +1177,9 @@ const STATIC_VTABLE: Vtable = Vtable {
     drop: static_drop,
 };
 
-unsafe fn static_clone(_: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Bytes {
-    let slice = slice::from_raw_parts(ptr, len);
-    Bytes::from_static(slice)
+unsafe fn static_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Bytes {
+    let data = AtomicPtr::new(data.load(Ordering::Relaxed));
+    Bytes::with_vtable(ptr, len, data, &STATIC_VTABLE)
 }
 
 unsafe fn static_to_vec(_: *mut (), ptr: *const u8, len: usize) -> Vec<u8> {

src/bytes_mut.rs

@@ -947,9 +947,8 @@ impl BytesMut {
     /// If the two `BytesMut` objects were previously contiguous and not mutated
     /// in a way that causes re-allocation i.e., if `other` was created by
     /// calling `split_off` on this `BytesMut`, then this is an `O(1)` operation
-    /// that just decreases a reference count and sets a few indices.
-    /// Otherwise this method degenerates to
-    /// `self.extend_from_slice(other.as_ref())`.
+    /// that just decreases a reference count and sets a few indices. Otherwise,
+    /// this method degenerates to `self.extend_from_slice(other.as_ref())`.
     ///
     /// # Examples
     ///
@@ -1088,13 +1087,11 @@ impl BytesMut {
             return Ok(());
         }
 
+        // Check if this block is right next to the other block
         let ptr = unsafe { self.ptr.as_ptr().add(self.len) };
-        if ptr == other.ptr.as_ptr()
-            && self.kind() == KIND_ARC
-            && other.kind() == KIND_ARC
-            && self.data == other.data
-        {
-            // Contiguous blocks, just combine directly
+        // ... and from the same backing data
+        if ptr == other.ptr.as_ptr() && self.data == other.data {
+            // ... then combine two blocks into one
             self.len += other.len;
             self.cap += other.cap;
             Ok(())

tests/test_bytes.rs

@@ -827,26 +827,198 @@ fn stress() {
 }
 
 #[test]
-fn partial_eq_bytesmut() {
+fn partial_eq_bytes_mut() {
     let bytes = Bytes::from(&b"The quick red fox"[..]);
-    let bytesmut = BytesMut::from(&b"The quick red fox"[..]);
-    assert!(bytes == bytesmut);
-    assert!(bytesmut == bytes);
+    let bytes_mut = BytesMut::from(&b"The quick red fox"[..]);
+    assert_eq!(bytes, bytes_mut);
+    assert_eq!(bytes_mut, bytes);
     let bytes2 = Bytes::from(&b"Jumped over the lazy brown dog"[..]);
-    assert!(bytes2 != bytesmut);
-    assert!(bytesmut != bytes2);
+    assert_ne!(bytes2, bytes_mut);
+    assert_ne!(bytes_mut, bytes2);
+}
+
+#[test]
+fn bytes_unsplit_split_static() {
+    let mut buf = Bytes::from_static(b"aaabbbcccddd");
+
+    let split = buf.split_off(6);
+    assert_eq!(b"aaabbb", &buf[..]);
+    assert_eq!(b"cccddd", &split[..]);
+
+    let new = buf.unsplit(split);
+    assert_eq!(b"aaabbbcccddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_split_mut() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaabbbcccddd");
+    let mut buf = buf.freeze();
+
+    let split = buf.split_off(6);
+    assert_eq!(b"aaabbb", &buf[..]);
+    assert_eq!(b"cccddd", &split[..]);
+
+    let new = buf.unsplit(split);
+    assert_eq!(b"aaabbbcccddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_empty_other() {
+    let buf = Bytes::from_static(b"aaabbbcccddd");
+
+    // empty other
+    let other = Bytes::new();
+
+    let new = buf.unsplit(other);
+    assert_eq!(b"aaabbbcccddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_empty_self() {
+    // empty self
+    let buf = Bytes::new();
+
+    let other = Bytes::from_static(b"aaabbbcccddd");
+
+    let new = buf.unsplit(other);
+    assert_eq!(b"aaabbbcccddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_two_static() {
+    let buf = Bytes::from_static(b"aaaabbbb");
+
+    let buf2 = Bytes::from_static(b"ccccdddd");
+
+    let new = buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_mut_and_static() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbb");
+    let buf = buf.freeze();
+
+    let buf2 = Bytes::from_static(b"ccccdddd");
+
+    let new = buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_two_mut() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbb");
+    let buf = buf.freeze();
+
+    let mut buf2 = BytesMut::with_capacity(64);
+    buf2.extend_from_slice(b"ccccdddd");
+    let buf2 = buf2.freeze();
+
+    let new = buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_two_static_split() {
+    let mut buf = Bytes::from_static(b"aaaabbbbeeee");
+
+    let _ = buf.split_off(8);
+
+    let mut buf2 = Bytes::from_static(b"ccccddddeeee");
+
+    let _ = buf2.split_off(8);
+
+    let new = buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_mut_and_static_split() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbbeeee");
+    let mut buf = buf.freeze();
+
+    let _ = buf.split_off(8);
+
+    let mut buf2 = Bytes::from_static(b"ccccddddeeee");
+
+    let _ = buf2.split_off(8);
+
+    let new = buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_two_mut_split() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbbeeee");
+    let mut buf = buf.freeze();
+
+    let _ = buf.split_off(8);
+
+    let mut buf2 = BytesMut::with_capacity(64);
+    buf2.extend_from_slice(b"ccccddddeeee");
+    let mut buf2 = buf2.freeze();
+
+    let _ = buf2.split_off(8);
+
+    let new = buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_static_non_contiguous() {
+    let mut buf = Bytes::from_static(b"aaaabbbbeeeeccccdddd");
+
+    let mut buf2 = buf.split_off(8);
+
+    let buf3 = buf2.split_off(4);
+
+    drop(buf2);
+
+    let new = buf.unsplit(buf3);
+    assert_eq!(b"aaaabbbbccccdddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_mut_non_contiguous() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbbeeeeccccdddd");
+    let mut buf = buf.freeze();
+
+    let mut buf2 = buf.split_off(8);
+
+    let buf3 = buf2.split_off(4);
+
+    drop(buf2);
+
+    let new = buf.unsplit(buf3);
+    assert_eq!(b"aaaabbbbccccdddd", &new[..]);
+}
+
+#[test]
+fn bytes_unsplit_from_split_static() {
+    let data: &[u8] = b"foobar";
+    let (a, b) = data.split_at(3);
+    let a = Bytes::from_static(a);
+    let b = Bytes::from_static(b);
+    let ab = a.unsplit(b);
+    assert_eq!(b"foobar", &ab[..]);
 }
 
 #[test]
 fn bytes_mut_unsplit_basic() {
     let mut buf = BytesMut::with_capacity(64);
     buf.extend_from_slice(b"aaabbbcccddd");
 
-    let splitted = buf.split_off(6);
+    let split = buf.split_off(6);
     assert_eq!(b"aaabbb", &buf[..]);
-    assert_eq!(b"cccddd", &splitted[..]);
+    assert_eq!(b"cccddd", &split[..]);
 
-    buf.unsplit(splitted);
+    buf.unsplit(split);
     assert_eq!(b"aaabbbcccddd", &buf[..]);
 }