diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index ae808faca..a792718af 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -1216,6 +1216,25 @@ impl BytesMut { slice::from_raw_parts_mut(ptr.cast(), len) } } + + /// 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`. + #[inline(always)] + unsafe fn extend_vec>(&mut self, iter: I) { + debug_assert_eq!(self.kind(), KIND_VEC); + + let off = self.get_vec_pos(); + 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; + } } impl Drop for BytesMut { @@ -1486,13 +1505,15 @@ 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); - // TODO: optimize - // 1. If self.kind() == KIND_VEC, use Vec::extend + if self.kind() == KIND_VEC { + unsafe { self.extend_vec(&mut iter) }; + } + for b in iter { self.put_u8(b); } diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 5d01b5b82..5b5fbbf15 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -622,6 +622,19 @@ 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] fn extend_from_slice_mut() { for &i in &[3, 34] {