diff --git a/src/bytes.rs b/src/bytes.rs index e71d24bb3e7c82809794427f414a5bc8997cc428..297659ddd10ac729f524f3f56edc6695b0d5e42b 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -648,6 +648,48 @@ impl Bytes { Err(self) } } + + /// Append given bytes to this object. + /// + /// If this `Bytes` object has not enough capacity, it is resized first. + /// It `Bytes` is shared (`refcount > 1`), it is copied first. + /// + /// This operation can be less effective than similar operation on `BytesMut`, + /// especially on small additions. + /// + /// # Examples + /// + /// ``` + /// use bytes::Bytes; + /// + /// let mut buf = Bytes::from("aabb"); + /// buf.extend_from_slice(b"ccdd"); + /// buf.extend_from_slice(b"eeff"); + /// + /// assert_eq!(b"aabbccddeeff", &buf[..]); + /// ``` + pub fn extend_from_slice(&mut self, extend: &[u8]) { + if extend.is_empty() { + return; + } + + let new_cap = self.len().checked_add(extend.len()).expect("capacity overflow"); + + let result = match mem::replace(self, Bytes::new()).try_mut() { + Ok(mut bytes_mut) => { + bytes_mut.extend_from_slice(extend); + bytes_mut + }, + Err(bytes) => { + let mut bytes_mut = BytesMut::with_capacity(new_cap); + bytes_mut.put_slice(&bytes); + bytes_mut.put_slice(extend); + bytes_mut + } + }; + + mem::replace(self, result.freeze()); + } } impl IntoBuf for Bytes { @@ -786,6 +828,38 @@ impl<'a> IntoIterator for &'a Bytes { } } +impl Extend<u8> for Bytes { + fn extend<T>(&mut self, iter: T) where T: IntoIterator<Item = u8> { + let iter = iter.into_iter(); + + let (lower, upper) = iter.size_hint(); + + // Avoid possible conversion into mut if there's nothing to add + if let Some(0) = upper { + return; + } + + let mut bytes_mut = match mem::replace(self, Bytes::new()).try_mut() { + Ok(bytes_mut) => bytes_mut, + Err(bytes) => { + let mut bytes_mut = BytesMut::with_capacity(bytes.len() + lower); + bytes_mut.put_slice(&bytes); + bytes_mut + } + }; + + bytes_mut.extend(iter); + + mem::replace(self, bytes_mut.freeze()); + } +} + +impl<'a> Extend<&'a u8> for Bytes { + fn extend<T>(&mut self, iter: T) where T: IntoIterator<Item = &'a u8> { + self.extend(iter.into_iter().map(|b| *b)) + } +} + /* * * ===== BytesMut ===== diff --git a/tests/test_bytes.rs b/tests/test_bytes.rs index 1a3435f2ef13b03bdeec544e5d8ffdf43c4a69ac..5c30475678611743a74bc38c657ea2984cfab845 100644 --- a/tests/test_bytes.rs +++ b/tests/test_bytes.rs @@ -357,12 +357,39 @@ fn inline_storage() { } #[test] -fn extend() { +fn extend_mut() { let mut bytes = BytesMut::with_capacity(0); bytes.extend(LONG); assert_eq!(*bytes, LONG[..]); } +#[test] +fn extend_shr() { + let mut bytes = Bytes::new(); + bytes.extend(LONG); + assert_eq!(*bytes, LONG[..]); +} + +#[test] +fn extend_from_slice_mut() { + for &i in &[3, 34] { + let mut bytes = BytesMut::new(); + bytes.extend_from_slice(&LONG[..i]); + bytes.extend_from_slice(&LONG[i..]); + assert_eq!(LONG[..], *bytes); + } +} + +#[test] +fn extend_from_slice_shr() { + for &i in &[3, 34] { + let mut bytes = Bytes::new(); + bytes.extend_from_slice(&LONG[..i]); + bytes.extend_from_slice(&LONG[i..]); + assert_eq!(LONG[..], *bytes); + } +} + #[test] fn from_static() { let mut a = Bytes::from_static(b"ab");