From ac956c1c136c50d0d1d42a7725d2061b76070012 Mon Sep 17 00:00:00 2001
From: Anthony Deschamps <anthony.j.deschamps@gmail.com>
Date: Thu, 24 Mar 2022 22:11:38 -0400
Subject: [PATCH] Implement Bytes::unsplit

This is based largely on the example of `BytesMut::unsplit`. If two
`Bytes` are contiguous and point to the same allocation, then they are
cheaply merged. Otherwise, a new `BytesMut` is allocated, copies data
from `self` and `other`, and `self` is replaced with `new.freeze()`.

Closes https://github.com/tokio-rs/bytes/issues/503
---
 src/bytes.rs | 60 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/src/bytes.rs b/src/bytes.rs
index f3e4c5cba..6ebac7eda 100644
--- a/src/bytes.rs
+++ b/src/bytes.rs
@@ -422,6 +422,48 @@ impl Bytes {
         ret
     }
 
+    /// 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` or `split_to` on
+    /// the same `BytesMut`, then this is an `O(1)` operation that
+    /// just decreases a reference count and sets a few
+    /// indices. Otherwise this method creates a new allocation and
+    /// copies both buffers into it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use bytes::BytesMut;
+    ///
+    /// let mut buf = BytesMut::with_capacity(64);
+    /// buf.extend_from_slice(b"aaabbbcccddd");
+    ///
+    /// let mut split_1 = buf.split_to(3).freeze();
+    /// assert_eq!(b"aaa", &split_1[..]);
+    /// assert_eq!(b"bbbcccddd", &buf[..]);
+    ///
+    /// let split_2 = buf.split_to(3).freeze();
+    /// assert_eq!(b"bbb", &split_2[..]);
+    /// assert_eq!(b"cccddd", &buf[..]);
+    ///
+    /// split_1.unsplit(split_2);
+    /// assert_eq!(b"aaabbb", &split_1[..]);
+    /// ```
+    pub fn unsplit(&mut self, other: Bytes) {
+        if self.is_empty() {
+            *self = other;
+            return;
+        }
+
+        if let Err(other) = self.try_unsplit(other) {
+            let mut new = crate::BytesMut::with_capacity(self.len + other.len);
+            new.extend_from_slice(self.as_ref());
+            new.extend_from_slice(other.as_ref());
+            *self = new.freeze();
+        }
+    }
+
     /// Shortens the buffer, keeping the first `len` bytes and dropping the
     /// rest.
     ///
@@ -503,6 +545,24 @@ impl Bytes {
         self.len -= by;
         self.ptr = self.ptr.add(by);
     }
+
+    fn try_unsplit(&mut self, other: Bytes) -> Result<(), Bytes> {
+        if other.len == 0 {
+            return Ok(());
+        }
+
+        let shared = self.data.load(Ordering::Acquire);
+        let other_shared = other.data.load(Ordering::Acquire);
+
+        let ptr = unsafe { self.ptr.add(self.len) };
+        if ptr == other.ptr && shared as usize & KIND_MASK == KIND_ARC && shared == other_shared {
+            // Contiguous blocks, just combine directly
+            self.len += other.len;
+            Ok(())
+        } else {
+            Err(other)
+        }
+    }
 }
 
 // Vtable must enforce this behavior