diff --git a/src/buf/buf_impl.rs b/src/buf/buf_impl.rs index c26681fb0..04c2f7c40 100644 --- a/src/buf/buf_impl.rs +++ b/src/buf/buf_impl.rs @@ -795,20 +795,28 @@ pub trait Buf { f64::from_bits(Self::get_u64_le(self)) } - /// Consumes remaining bytes inside self and returns new instance of `Bytes` + /// Consumes `len` bytes inside self and returns new instance of `Bytes` + /// with this data. + /// + /// This function may be optimized by the underlying type to avoid actual + /// copies. For example, `Bytes` implementation will do a shallow copy + /// (ref-count increment). /// /// # Examples /// /// ``` /// use bytes::Buf; /// - /// let bytes = (&b"hello world"[..]).to_bytes(); - /// assert_eq!(&bytes[..], &b"hello world"[..]); + /// let bytes = (&b"hello world"[..]).copy_to_bytes(5); + /// assert_eq!(&bytes[..], &b"hello"[..]); /// ``` - fn to_bytes(&mut self) -> crate::Bytes { + fn copy_to_bytes(&mut self, len: usize) -> crate::Bytes { use super::BufMut; - let mut ret = crate::BytesMut::with_capacity(self.remaining()); - ret.put(self); + + assert!(len <= self.remaining(), "`len` greater than remaining"); + + let mut ret = crate::BytesMut::with_capacity(len); + ret.put(self.take(len)); ret.freeze() } @@ -852,7 +860,7 @@ pub trait Buf { /// /// let mut chain = b"hello "[..].chain(&b"world"[..]); /// - /// let full = chain.to_bytes(); + /// let full = chain.copy_to_bytes(11); /// assert_eq!(full.bytes(), b"hello world"); /// ``` fn chain(self, next: U) -> Chain @@ -993,8 +1001,8 @@ macro_rules! deref_forward_buf { (**self).get_int_le(nbytes) } - fn to_bytes(&mut self) -> crate::Bytes { - (**self).to_bytes() + fn copy_to_bytes(&mut self, len: usize) -> crate::Bytes { + (**self).copy_to_bytes(len) } }; } diff --git a/src/buf/chain.rs b/src/buf/chain.rs index db6219281..c21fd35ae 100644 --- a/src/buf/chain.rs +++ b/src/buf/chain.rs @@ -21,7 +21,7 @@ use std::io::IoSlice; /// let mut buf = (&b"hello "[..]) /// .chain(&b"world"[..]); /// -/// let full: Bytes = buf.to_bytes(); +/// let full: Bytes = buf.copy_to_bytes(11); /// assert_eq!(full[..], b"hello world"[..]); /// ``` /// @@ -68,7 +68,7 @@ impl Chain { /// /// buf.first_mut().advance(1); /// - /// let full = buf.to_bytes(); + /// let full = buf.copy_to_bytes(9); /// assert_eq!(full, b"elloworld"[..]); /// ``` pub fn first_mut(&mut self) -> &mut T { @@ -103,7 +103,7 @@ impl Chain { /// /// buf.last_mut().advance(1); /// - /// let full = buf.to_bytes(); + /// let full = buf.copy_to_bytes(10); /// assert_eq!(full, b"hello orld"[..]); /// ``` pub fn last_mut(&mut self) -> &mut U { diff --git a/src/bytes.rs b/src/bytes.rs index 66b941d99..867f86bb8 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -548,8 +548,14 @@ impl Buf for Bytes { } } - fn to_bytes(&mut self) -> crate::Bytes { - core::mem::replace(self, Bytes::new()) + fn copy_to_bytes(&mut self, len: usize) -> crate::Bytes { + if len == self.remaining() { + core::mem::replace(self, Bytes::new()) + } else { + let ret = self.slice(..len); + self.advance(len); + ret + } } } diff --git a/src/bytes_mut.rs b/src/bytes_mut.rs index 38f1ed53d..c3abadf14 100644 --- a/src/bytes_mut.rs +++ b/src/bytes_mut.rs @@ -961,8 +961,8 @@ impl Buf for BytesMut { } } - fn to_bytes(&mut self) -> crate::Bytes { - self.split().freeze() + fn copy_to_bytes(&mut self, len: usize) -> crate::Bytes { + self.split_to(len).freeze() } } diff --git a/tests/test_buf.rs b/tests/test_buf.rs index 17bdd54e8..bc1287323 100644 --- a/tests/test_buf.rs +++ b/tests/test_buf.rs @@ -101,3 +101,20 @@ fn test_deref_buf_forwards() { assert_eq!((Box::new(Special) as Box).get_u8(), b'x'); assert_eq!(Box::new(Special).get_u8(), b'x'); } + +#[test] +fn copy_to_bytes_less() { + let mut buf = &b"hello world"[..]; + + let bytes = buf.copy_to_bytes(5); + assert_eq!(bytes, &b"hello"[..]); + assert_eq!(buf, &b" world"[..]) +} + +#[test] +#[should_panic] +fn copy_to_bytes_overflow() { + let mut buf = &b"hello world"[..]; + + let _bytes = buf.copy_to_bytes(12); +} diff --git a/tests/test_chain.rs b/tests/test_chain.rs index 9be434f37..0f362846e 100644 --- a/tests/test_chain.rs +++ b/tests/test_chain.rs @@ -9,7 +9,7 @@ fn collect_two_bufs() { let a = Bytes::from(&b"hello"[..]); let b = Bytes::from(&b"world"[..]); - let res = a.chain(b).to_bytes(); + let res = a.chain(b).copy_to_bytes(10); assert_eq!(res, &b"helloworld"[..]); }