From 349b8233f850272bc6f1a4c388abbb37c7d60c91 Mon Sep 17 00:00:00 2001 From: Mai Thanh Minh Date: Tue, 27 Jan 2026 09:37:37 +0700 Subject: [PATCH] Implement `Bytes::unsplit` and `Bytes::try_unsplit` --- src/bytes.rs | 117 +++++++++++++++++++++++++-- src/bytes_mut.rs | 15 ++-- tests/test_bytes.rs | 190 +++++++++++++++++++++++++++++++++++++++++--- 3 files changed, 298 insertions(+), 24 deletions(-) diff --git a/src/bytes.rs b/src/bytes.rs index aed4d7cdb..113f3c00e 100644 --- a/src/bytes.rs +++ b/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 { diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index ae808faca..7b685ab63 100644 --- a/src/bytes_mut.rs +++ b/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(()) diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 5d01b5b82..988b59dac 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -827,14 +827,186 @@ 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] @@ -842,11 +1014,11 @@ 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[..]); }