diff --git a/tokio-buf/src/util/from.rs b/tokio-buf/src/util/from.rs index b94ebc2d8..1c09e6412 100644 --- a/tokio-buf/src/util/from.rs +++ b/tokio-buf/src/util/from.rs @@ -1,6 +1,6 @@ use SizeHint; -use bytes::{Buf, BufMut}; +use bytes::{Buf, BufMut, Bytes}; use std::usize; @@ -47,12 +47,18 @@ pub struct CollectVecError { _p: (), } +/// Error returned from collecting into a `Bytes` +#[derive(Debug)] +pub struct CollectBytesError { + _p: (), +} + impl FromBufStream for Vec { type Builder = Vec; type Error = CollectVecError; - fn builder(_hint: &SizeHint) -> Vec { - Vec::new() + fn builder(hint: &SizeHint) -> Vec { + Vec::with_capacity(hint.lower() as usize) } fn extend(builder: &mut Self, buf: &mut T, hint: &SizeHint) -> Result<(), Self::Error> { @@ -110,3 +116,21 @@ impl FromBufStream for Vec { Ok(builder) } } + +impl FromBufStream for Bytes { + type Builder = Vec; + type Error = CollectBytesError; + + fn builder(hint: &SizeHint) -> Vec { + as FromBufStream>::builder(hint) + } + + fn extend(builder: &mut Vec, buf: &mut T, hint: &SizeHint) -> Result<(), Self::Error> { + as FromBufStream>::extend(builder, buf, hint) + .map_err(|_| CollectBytesError { _p: () }) + } + + fn build(builder: Vec) -> Result { + Ok(builder.into()) + } +} diff --git a/tokio-buf/src/util/mod.rs b/tokio-buf/src/util/mod.rs index 2a3dfbf7f..c9a0b64ee 100644 --- a/tokio-buf/src/util/mod.rs +++ b/tokio-buf/src/util/mod.rs @@ -14,7 +14,7 @@ pub mod error { //! Error types pub use super::collect::CollectError; - pub use super::from::CollectVecError; + pub use super::from::{CollectBytesError, CollectVecError}; pub use super::limit::LimitError; } diff --git a/tokio-buf/tests/buf_stream_ext.rs b/tokio-buf/tests/buf_stream_ext.rs index 92658ee1a..cda2aa95b 100644 --- a/tokio-buf/tests/buf_stream_ext.rs +++ b/tokio-buf/tests/buf_stream_ext.rs @@ -1,10 +1,10 @@ -#![cfg(feature = "ext")] +#![cfg(feature = "util")] extern crate bytes; extern crate futures; extern crate tokio_buf; -use bytes::Buf; +use bytes::{Buf, Bytes}; use futures::Async::*; use futures::Future; use tokio_buf::{BufStream, BufStreamExt}; @@ -48,47 +48,58 @@ fn chain() { // ===== Test `collect()` ===== +macro_rules! test_collect_impl { + ($t:ty $(, $capacity:ident)*) => { + // While unfortunate, this test makes some assumptions on vec's resizing + // behavior. + // + // Collect one + // + let bs = one("hello world"); + + let vec: $t = bs.collect().wait().unwrap(); + + assert_eq!(vec, &b"hello world"[..]); + $( assert_eq!(vec.$capacity(), 64); )* + + // Collect one, with size hint + // + let mut bs = one("hello world"); + bs.size_hint.set_lower(11); + + let vec: $t = bs.collect().wait().unwrap(); + + assert_eq!(vec, &b"hello world"[..]); + $( assert_eq!(vec.$capacity(), 64); )* + + // Collect one, with size hint + // + let mut bs = one("hello world"); + bs.size_hint.set_lower(10); + + let vec: $t = bs.collect().wait().unwrap(); + + assert_eq!(vec, &b"hello world"[..]); + $( assert_eq!(vec.$capacity(), 64); )* + + // Collect many + // + let bs = list(&["hello", " ", "world", ", one two three"]); + + let vec: $t = bs.collect().wait().unwrap(); + + assert_eq!(vec, &b"hello world, one two three"[..]); + } +} + #[test] fn collect_vec() { - // While unfortunate, this test makes some assumptions on vec's resizing - // behavior. - // - // Collect one - // - let bs = one("hello world"); + test_collect_impl!(Vec, capacity); +} - let vec: Vec = bs.collect().wait().unwrap(); - - assert_eq!(vec, b"hello world"); - assert_eq!(vec.capacity(), 64); - - // Collect one, with size hint - // - let mut bs = one("hello world"); - bs.size_hint.set_lower(11); - - let vec: Vec = bs.collect().wait().unwrap(); - - assert_eq!(vec, b"hello world"); - assert_eq!(vec.capacity(), 64); - - // Collect one, with size hint - // - let mut bs = one("hello world"); - bs.size_hint.set_lower(10); - - let vec: Vec = bs.collect().wait().unwrap(); - - assert_eq!(vec, b"hello world"); - assert_eq!(vec.capacity(), 64); - - // Collect many - // - let bs = list(&["hello", " ", "world", ", one two three"]); - - let vec: Vec = bs.collect().wait().unwrap(); - - assert_eq!(vec, b"hello world, one two three"); +#[test] +fn collect_bytes() { + test_collect_impl!(Bytes); } // ===== Test limit() =====