buf: implement FromBufStream for Bytes (#1009)

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<T: Buf> FromBufStream<T> for Vec<u8> {
     type Builder = Vec<u8>;
     type Error = CollectVecError;
 
-    fn builder(_hint: &SizeHint) -> Vec<u8> {
+    fn builder(hint: &SizeHint) -> Vec<u8> {
-        Vec::new()
+        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<T: Buf> FromBufStream<T> for Vec<u8> {
         Ok(builder)
     }
 }
+
+impl<T: Buf> FromBufStream<T> for Bytes {
+    type Builder = Vec<u8>;
+    type Error = CollectBytesError;
+
+    fn builder(hint: &SizeHint) -> Vec<u8> {
+        <Vec<u8> as FromBufStream<T>>::builder(hint)
+    }
+
+    fn extend(builder: &mut Vec<u8>, buf: &mut T, hint: &SizeHint) -> Result<(), Self::Error> {
+        <Vec<u8> as FromBufStream<T>>::extend(builder, buf, hint)
+            .map_err(|_| CollectBytesError { _p: () })
+    }
+
+    fn build(builder: Vec<u8>) -> Result<Self, Self::Error> {
+        Ok(builder.into())
+    }
+}

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;
 }
 

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
+    test_collect_impl!(Vec<u8>, capacity);
-    // behavior.
+}
-    //
-    // Collect one
-    //
-    let bs = one("hello world");
 
-    let vec: Vec<u8> = bs.collect().wait().unwrap();
+#[test]
-
+fn collect_bytes() {
-    assert_eq!(vec, b"hello world");
+    test_collect_impl!(Bytes);
-    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<u8> = 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<u8> = 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<u8> = bs.collect().wait().unwrap();
-
-    assert_eq!(vec, b"hello world, one two three");
 }
 
 // ===== Test limit() =====