From 7ff78e7e8831b8ef61697016cf5a91774e7b1064 Mon Sep 17 00:00:00 2001
From: jq-rs <juhamatk@gmail.com>
Date: Wed, 3 Jan 2018 19:54:51 +0200
Subject: [PATCH] Add support for unsplit() to BytesMut (#162)

Add support for unsplit() to BytesMut which combines splitted contiguous memory blocks efficiently.
---
 src/bytes.rs        | 40 ++++++++++++++++++
 tests/test_bytes.rs | 98 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 138 insertions(+)

diff --git a/src/bytes.rs b/src/bytes.rs
index dea63e2..5dba49e 100644
--- a/src/bytes.rs
+++ b/src/bytes.rs
@@ -1378,6 +1378,46 @@ impl BytesMut {
         self.reserve(extend.len());
         self.put_slice(extend);
     }
+
+    /// Combine splitted BytesMut objects back as contiguous.
+    ///
+    /// If `BytesMut` objects were not contiguous originally, they will be extended.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::BytesMut;
+    ///
+    /// let mut buf = BytesMut::with_capacity(64);
+    /// buf.extend_from_slice(b"aaabbbcccddd");
+    ///
+    /// let splitted = buf.split_off(6);
+    /// assert_eq!(b"aaabbb", &buf[..]);
+    /// assert_eq!(b"cccddd", &splitted[..]);
+    ///
+    /// buf.unsplit(splitted);
+    /// assert_eq!(b"aaabbbcccddd", &buf[..]);
+    /// ```
+    pub fn unsplit(&mut self, other: BytesMut) {
+        let ptr;
+
+        unsafe {
+            ptr = self.inner.ptr.offset(self.inner.len as isize); 
+        }
+        if ptr == other.inner.ptr &&
+           self.inner.kind() == KIND_ARC &&
+           other.inner.kind() == KIND_ARC
+        {
+            debug_assert_eq!(self.inner.arc.load(Acquire),
+                             other.inner.arc.load(Acquire));
+            // Contiguous blocks, just combine directly
+            self.inner.len += other.inner.len;
+            self.inner.cap += other.inner.cap;
+        }
+        else {
+            self.extend(other);
+        }
+    }
 }
 
 impl BufMut for BytesMut {
diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs
index 35ba8ac..e5b1fe7 100644
--- a/tests/test_bytes.rs
+++ b/tests/test_bytes.rs
@@ -552,3 +552,101 @@ fn partial_eq_bytesmut() {
     assert!(bytes2 != bytesmut);
     assert!(bytesmut != bytes2);
 }
+
+#[test]
+fn unsplit_basic() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaabbbcccddd");
+
+    let splitted = buf.split_off(6);
+    assert_eq!(b"aaabbb", &buf[..]);
+    assert_eq!(b"cccddd", &splitted[..]);
+
+    buf.unsplit(splitted);
+    assert_eq!(b"aaabbbcccddd", &buf[..]);
+}
+
+#[test]
+fn unsplit_inline_arc() {
+    let mut buf = BytesMut::with_capacity(8); //inline
+    buf.extend_from_slice(b"aaaabbbb");
+
+    let mut buf2 = BytesMut::with_capacity(64);
+    buf2.extend_from_slice(b"ccccddddeeee");
+
+    buf2.split_off(8); //arc
+
+    buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &buf[..]);
+}
+
+#[test]
+fn unsplit_arc_inline() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbbeeee");
+
+    buf.split_off(8); //arc
+
+    let mut buf2 = BytesMut::with_capacity(8); //inline
+    buf2.extend_from_slice(b"ccccdddd");
+
+    buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &buf[..]);
+
+}
+
+#[test]
+fn unsplit_both_inline() {
+    let mut buf = BytesMut::with_capacity(16); //inline
+    buf.extend_from_slice(b"aaaabbbbccccdddd");
+
+    let splitted = buf.split_off(8); // both inline
+    assert_eq!(b"aaaabbbb", &buf[..]);
+    assert_eq!(b"ccccdddd", &splitted[..]);
+
+    buf.unsplit(splitted);
+    assert_eq!(b"aaaabbbbccccdddd", &buf[..]);
+}
+
+
+#[test]
+fn unsplit_arc_different() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbbeeee");
+
+    buf.split_off(8); //arc
+
+    let mut buf2 = BytesMut::with_capacity(64);
+    buf2.extend_from_slice(b"ccccddddeeee");
+
+    buf2.split_off(8); //arc
+
+    buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &buf[..]);
+}
+
+#[test]
+fn unsplit_arc_non_contiguous() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbbeeeeccccdddd");
+
+    let mut buf2 = buf.split_off(8); //arc
+
+    let buf3 = buf2.split_off(4); //arc
+
+    buf.unsplit(buf3);
+    assert_eq!(b"aaaabbbbccccdddd", &buf[..]);
+}
+
+#[test]
+fn unsplit_two_split_offs() {
+    let mut buf = BytesMut::with_capacity(64);
+    buf.extend_from_slice(b"aaaabbbbccccdddd");
+
+    let mut buf2 = buf.split_off(8); //arc
+    let buf3 = buf2.split_off(4); //arc
+
+    buf2.unsplit(buf3);
+    buf.unsplit(buf2);
+    assert_eq!(b"aaaabbbbccccdddd", &buf[..]);
+}