Skip to content

Commit

Permalink
Add Buf::copy_to_bytes(len) (#439)
Browse files Browse the repository at this point in the history
This method replaces `Buf::to_bytes()`, providing a method that copies a
subset of the remaining buffer into a `Bytes` value. As this is strictly
more flexible, `to_bytes()` is removed.

Fixes: #129, #398
  • Loading branch information
carllerche authored Oct 20, 2020
1 parent 5866839 commit 39de065
Show file tree
Hide file tree
Showing 6 changed files with 48 additions and 17 deletions.
26 changes: 17 additions & 9 deletions src/buf/buf_impl.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}

Expand Down Expand Up @@ -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<U: Buf>(self, next: U) -> Chain<Self, U>
Expand Down Expand Up @@ -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)
}
};
}
Expand Down
6 changes: 3 additions & 3 deletions src/buf/chain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"[..]);
/// ```
///
Expand Down Expand Up @@ -68,7 +68,7 @@ impl<T, U> Chain<T, U> {
///
/// 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 {
Expand Down Expand Up @@ -103,7 +103,7 @@ impl<T, U> Chain<T, U> {
///
/// 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 {
Expand Down
10 changes: 8 additions & 2 deletions src/bytes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
}
}

Expand Down
4 changes: 2 additions & 2 deletions src/bytes_mut.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}
}

Expand Down
17 changes: 17 additions & 0 deletions tests/test_buf.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,3 +101,20 @@ fn test_deref_buf_forwards() {
assert_eq!((Box::new(Special) as Box<dyn Buf>).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);
}
2 changes: 1 addition & 1 deletion tests/test_chain.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"[..]);
}

Expand Down

0 comments on commit 39de065

Please sign in to comment.