From fe9fa352374157b29f0c55e7e3f1b949c552cf85 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Mon, 15 Jun 2026 17:49:37 +0100 Subject: [PATCH 1/2] Specialize BytesMut::extend when Vec backed Signed-off-by: Adam Gutglick --- src/bytes_mut.rs | 50 +++++++++++++++++++++++++++++++++++---- tests/test_bytes.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 103 insertions(+), 4 deletions(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index ae808faca..4db6bf34f 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1216,6 +1216,26 @@ impl BytesMut { slice::from_raw_parts_mut(ptr.cast(), len) } } + + /// Extends the buffer from an `IntoIterator`. + /// + /// # SAFETY + /// + /// The caller must ensure that `self.kind()` is `KIND_VEC`. + unsafe fn extend_vec>(&mut self, iter: I) { + debug_assert_eq!(self.kind(), KIND_VEC); + + let off = self.get_vec_pos(); + let vec = ManuallyDrop::new(rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off)); + + let mut guard = VecRebuildGuard { + bytes: self, + vec, + off, + }; + + guard.vec.extend(iter); + } } impl Drop for BytesMut { @@ -1491,10 +1511,12 @@ impl Extend for BytesMut { let (lower, _) = iter.size_hint(); self.reserve(lower); - // TODO: optimize - // 1. If self.kind() == KIND_VEC, use Vec::extend - for b in iter { - self.put_u8(b); + if self.kind() == KIND_VEC { + unsafe { self.extend_vec(iter) }; + } else { + for b in iter { + self.put_u8(b); + } } } } @@ -1887,6 +1909,26 @@ unsafe fn rebuild_vec(ptr: *mut u8, mut len: usize, mut cap: usize, off: usize) Vec::from_raw_parts(ptr, len, cap) } +/// RAII guard in case the extension panics and is caught before +/// we update the state. +/// +/// On drop - updates the `BytesMut` len, cap and ptr. +struct VecRebuildGuard<'a> { + bytes: &'a mut BytesMut, + vec: ManuallyDrop>, + off: usize, +} + +impl Drop for VecRebuildGuard<'_> { + fn drop(&mut self) { + unsafe { + self.bytes.len = self.vec.len() - self.off; + self.bytes.cap = self.vec.capacity() - self.off; + self.bytes.ptr = vptr(self.vec.as_mut_ptr().add(self.off)); + } + } +} + // ===== impl SharedVtable ===== static SHARED_VTABLE: Vtable = Vtable { diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 5d01b5b82..c0e8b7d54 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -622,6 +622,63 @@ fn extend_mut() { assert_eq!(*bytes, LONG[..]); } +#[test] +fn extend_mut_after_advance() { + let mut bytes = BytesMut::with_capacity(16); + bytes.extend_from_slice(b"abcdef"); + bytes.advance(2); + + bytes.extend([b'g', b'h']); + + assert_eq!(bytes.len(), 6); + assert_eq!(&bytes[..], b"cdefgh"); + assert_eq!(bytes.capacity(), 14); +} + +#[test] +#[cfg_attr(not(panic = "unwind"), ignore)] +fn extend_mut_keeps_buffer_valid_on_iterator_panic() { + struct PanicAfter { + next: u8, + remaining: usize, + } + + impl Iterator for PanicAfter { + type Item = u8; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + panic!("test-triggered iterator panic"); + } + + let next = self.next; + self.next += 1; + self.remaining -= 1; + Some(next) + } + + fn size_hint(&self) -> (usize, Option) { + (0, None) + } + } + + let mut bytes = BytesMut::with_capacity(16); + bytes.extend_from_slice(b"a"); + + let result = panic::catch_unwind(AssertUnwindSafe(|| { + bytes.extend(PanicAfter { + next: b'b', + remaining: 3, + }); + })); + + assert!(result.is_err()); + assert_eq!(&bytes[..], b"abcd"); + + bytes.extend_from_slice(b"ef"); + assert_eq!(&bytes[..], b"abcdef"); +} + #[test] fn extend_from_slice_mut() { for &i in &[3, 34] { From 8884eadbc7d20c449305a644b11dbc019b95c50d Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Wed, 16 Sep 2026 11:56:10 +0100 Subject: [PATCH 2/2] Simplify Signed-off-by: Adam Gutglick --- src/bytes_mut.rs | 53 ++++++++++++++------------------------------- tests/test_bytes.rs | 44 ------------------------------------- 2 files changed, 16 insertions(+), 81 deletions(-) diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 4db6bf34f..a792718af 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1217,24 +1217,23 @@ impl BytesMut { } } - /// Extends the buffer from an `IntoIterator`. + /// Appends bytes from `iter` into the spare capacity, without growing. + /// + /// Goes through `Vec::extend`, which carries std's `TrustedLen` + /// specialization (slice-like iterators become a memcpy). /// /// # SAFETY /// /// The caller must ensure that `self.kind()` is `KIND_VEC`. - unsafe fn extend_vec>(&mut self, iter: I) { + #[inline(always)] + unsafe fn extend_vec>(&mut self, iter: I) { debug_assert_eq!(self.kind(), KIND_VEC); let off = self.get_vec_pos(); - let vec = ManuallyDrop::new(rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off)); - - let mut guard = VecRebuildGuard { - bytes: self, - vec, - off, - }; - - guard.vec.extend(iter); + let spare = self.cap - self.len; + let mut vec = ManuallyDrop::new(rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off)); + vec.extend(iter.take(spare)); + self.len = vec.len() - off; } } @@ -1506,17 +1505,17 @@ impl Extend for BytesMut { where T: IntoIterator, { - let iter = iter.into_iter(); + let mut iter = iter.into_iter(); let (lower, _) = iter.size_hint(); self.reserve(lower); if self.kind() == KIND_VEC { - unsafe { self.extend_vec(iter) }; - } else { - for b in iter { - self.put_u8(b); - } + unsafe { self.extend_vec(&mut iter) }; + } + + for b in iter { + self.put_u8(b); } } } @@ -1909,26 +1908,6 @@ unsafe fn rebuild_vec(ptr: *mut u8, mut len: usize, mut cap: usize, off: usize) Vec::from_raw_parts(ptr, len, cap) } -/// RAII guard in case the extension panics and is caught before -/// we update the state. -/// -/// On drop - updates the `BytesMut` len, cap and ptr. -struct VecRebuildGuard<'a> { - bytes: &'a mut BytesMut, - vec: ManuallyDrop>, - off: usize, -} - -impl Drop for VecRebuildGuard<'_> { - fn drop(&mut self) { - unsafe { - self.bytes.len = self.vec.len() - self.off; - self.bytes.cap = self.vec.capacity() - self.off; - self.bytes.ptr = vptr(self.vec.as_mut_ptr().add(self.off)); - } - } -} - // ===== impl SharedVtable ===== static SHARED_VTABLE: Vtable = Vtable { diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index c0e8b7d54..5b5fbbf15 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -635,50 +635,6 @@ fn extend_mut_after_advance() { assert_eq!(bytes.capacity(), 14); } -#[test] -#[cfg_attr(not(panic = "unwind"), ignore)] -fn extend_mut_keeps_buffer_valid_on_iterator_panic() { - struct PanicAfter { - next: u8, - remaining: usize, - } - - impl Iterator for PanicAfter { - type Item = u8; - - fn next(&mut self) -> Option { - if self.remaining == 0 { - panic!("test-triggered iterator panic"); - } - - let next = self.next; - self.next += 1; - self.remaining -= 1; - Some(next) - } - - fn size_hint(&self) -> (usize, Option) { - (0, None) - } - } - - let mut bytes = BytesMut::with_capacity(16); - bytes.extend_from_slice(b"a"); - - let result = panic::catch_unwind(AssertUnwindSafe(|| { - bytes.extend(PanicAfter { - next: b'b', - remaining: 3, - }); - })); - - assert!(result.is_err()); - assert_eq!(&bytes[..], b"abcd"); - - bytes.extend_from_slice(b"ef"); - assert_eq!(&bytes[..], b"abcdef"); -} - #[test] fn extend_from_slice_mut() { for &i in &[3, 34] {