buf: misc polish (#924)

- Rename feature flag `util`.
- Rename module `util`
- Move `error` module into `util`.
- Move `BufStream` impls into dedicated file.

tokio-buf/src/util/chain.rs

@@ -0,0 +1,46 @@
+use BufStream;
+
+use either::Either;
+use futures::Poll;
+
+/// A buf stream that sequences two buf streams together.
+///
+/// `Chain` values are produced by the `chain` function on `BufStream`.
+#[derive(Debug)]
+pub struct Chain<T, U> {
+    left: Option<T>,
+    right: U,
+}
+
+impl<T, U> Chain<T, U> {
+    pub(crate) fn new(left: T, right: U) -> Chain<T, U> {
+        Chain {
+            left: Some(left),
+            right,
+        }
+    }
+}
+
+impl<T, U> BufStream for Chain<T, U>
+where
+    T: BufStream,
+    U: BufStream<Error = T::Error>,
+{
+    type Item = Either<T::Item, U::Item>;
+    type Error = T::Error;
+
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        if let Some(ref mut stream) = self.left {
+            let res = try_ready!(stream.poll_buf());
+
+            if res.is_some() {
+                return Ok(res.map(Either::Left).into());
+            }
+        }
+
+        self.left = None;
+
+        let res = try_ready!(self.right.poll_buf());
+        Ok(res.map(Either::Right).into())
+    }
+}

tokio-buf/src/util/collect.rs

@@ -0,0 +1,101 @@
+use super::FromBufStream;
+use BufStream;
+
+use futures::{Future, Poll};
+
+/// Consumes a buf stream, collecting the data into a single byte container.
+///
+/// `Collect` values are produced by `BufStream::collect`.
+#[derive(Debug)]
+pub struct Collect<T, U>
+where
+    T: BufStream,
+    U: FromBufStream<T::Item>,
+{
+    stream: T,
+    builder: Option<U::Builder>,
+}
+
+/// Errors returned from `Collect` future.
+#[derive(Debug)]
+pub struct CollectError<T, U> {
+    inner: Error<T, U>,
+}
+
+#[derive(Debug)]
+enum Error<T, U> {
+    Stream(T),
+    Collect(U),
+}
+
+impl<T, U> Collect<T, U>
+where
+    T: BufStream,
+    U: FromBufStream<T::Item>,
+{
+    pub(crate) fn new(stream: T) -> Collect<T, U> {
+        let builder = U::builder(&stream.size_hint());
+
+        Collect {
+            stream,
+            builder: Some(builder),
+        }
+    }
+}
+
+impl<T, U> Future for Collect<T, U>
+where
+    T: BufStream,
+    U: FromBufStream<T::Item>,
+{
+    type Item = U;
+    type Error = CollectError<T::Error, U::Error>;
+
+    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
+        loop {
+            let res = self.stream.poll_buf().map_err(|err| {
+                let inner = Error::Stream(err);
+                CollectError { inner }
+            });
+
+            match try_ready!(res) {
+                Some(mut buf) => {
+                    let builder = self.builder.as_mut().expect("cannot poll after done");
+
+                    U::extend(builder, &mut buf, &self.stream.size_hint()).map_err(|err| {
+                        let inner = Error::Collect(err);
+                        CollectError { inner }
+                    })?;
+                }
+                None => {
+                    let builder = self.builder.take().expect("cannot poll after done");
+                    let value = U::build(builder).map_err(|err| {
+                        let inner = Error::Collect(err);
+                        CollectError { inner }
+                    })?;
+                    return Ok(value.into());
+                }
+            }
+        }
+    }
+}
+
+// ===== impl CollectError =====
+
+impl<T, U> CollectError<T, U> {
+    /// Returns `true` if the error was caused by polling the stream.
+    pub fn is_stream_err(&self) -> bool {
+        match self.inner {
+            Error::Stream(_) => true,
+            _ => false,
+        }
+    }
+
+    /// Returns `true` if the error happened while collecting the data.
+    pub fn is_collect_err(&self) -> bool {
+        match self.inner {
+            Error::Collect(_) => true,
+            _ => false,
+        }
+    }
+}

tokio-buf/src/util/from.rs

@@ -0,0 +1,112 @@
+use SizeHint;
+
+use bytes::{Buf, BufMut};
+
+use std::usize;
+
+/// Conversion from a `BufStream`.
+///
+/// By implementing `FromBufStream` for a type, you define how it will be
+/// created from a buf stream. This is common for types which describe byte
+/// storage of some kind.
+///
+/// `FromBufStream` is rarely called explicitly, and it is instead used through
+/// `BufStream`'s `collect` method.
+pub trait FromBufStream<T: Buf>: Sized {
+    /// Type that is used to build `Self` while the `BufStream` is being
+    /// consumed.
+    type Builder;
+
+    /// Error that might happen on conversion.
+    type Error;
+
+    /// Create a new, empty, builder. The provided `hint` can be used to inform
+    /// reserving capacity.
+    fn builder(hint: &SizeHint) -> Self::Builder;
+
+    /// Extend the builder with the `Buf`.
+    ///
+    /// This method is called whenever a new `Buf` value is obtained from the
+    /// buf stream.
+    ///
+    /// The provided size hint represents the state of the stream **after**
+    /// `buf` has been yielded. The lower bound represents the minimum amount of
+    /// data that will be provided after this call to `extend` returns.
+    fn extend(builder: &mut Self::Builder, buf: &mut T, hint: &SizeHint)
+        -> Result<(), Self::Error>;
+
+    /// Finalize the building of `Self`.
+    ///
+    /// Called once the buf stream is fully consumed.
+    fn build(builder: Self::Builder) -> Result<Self, Self::Error>;
+}
+
+/// Error returned from collecting into a `Vec<u8>`
+#[derive(Debug)]
+pub struct CollectVecError {
+    _p: (),
+}
+
+impl<T: Buf> FromBufStream<T> for Vec<u8> {
+    type Builder = Vec<u8>;
+    type Error = CollectVecError;
+
+    fn builder(_hint: &SizeHint) -> Vec<u8> {
+        Vec::new()
+    }
+
+    fn extend(builder: &mut Self, buf: &mut T, hint: &SizeHint) -> Result<(), Self::Error> {
+        let lower = hint.lower();
+
+        // If the lower bound is greater than `usize::MAX` then we have a
+        // problem
+        if lower > usize::MAX as u64 {
+            return Err(CollectVecError { _p: () });
+        }
+
+        let mut reserve = lower as usize;
+
+        // If `upper` is set, use this value if it is less than or equal to 64.
+        // This only really impacts the first iteration.
+        match hint.upper() {
+            Some(upper) if upper <= 64 => {
+                reserve = upper as usize;
+            }
+            _ => {}
+        }
+
+        // hint.lower() represents the minimum amount of data that will be
+        // received *after* this function call. We reserve this amount on top of
+        // the amount of data in `buf`.
+        reserve = match reserve.checked_add(buf.remaining()) {
+            Some(n) => n,
+            None => return Err(CollectVecError { _p: () }),
+        };
+
+        // Always reserve 64 bytes the first time, unless `upper` is set and is
+        // less than 64.
+        if builder.is_empty() {
+            reserve = reserve.max(match hint.upper() {
+                Some(upper) if upper < 64 => upper as usize,
+                _ => 64,
+            });
+        }
+
+        // Make sure overflow won't happen when reserving
+        if reserve.checked_add(builder.len()).is_none() {
+            return Err(CollectVecError { _p: () });
+        }
+
+        // Reserve space
+        builder.reserve(reserve);
+
+        // Copy the data
+        builder.put(buf);
+
+        Ok(())
+    }
+
+    fn build(builder: Self) -> Result<Self, Self::Error> {
+        Ok(builder)
+    }
+}

tokio-buf/src/util/limit.rs

@@ -0,0 +1,76 @@
+use BufStream;
+
+use bytes::Buf;
+use futures::Poll;
+
+/// Limits the stream to a maximum amount of data.
+#[derive(Debug)]
+pub struct Limit<T> {
+    stream: T,
+    remaining: u64,
+}
+
+/// Errors returned from `Limit`.
+#[derive(Debug)]
+pub struct LimitError<T> {
+    /// When `None`, limit was reached
+    inner: Option<T>,
+}
+
+impl<T> Limit<T> {
+    pub(crate) fn new(stream: T, amount: u64) -> Limit<T> {
+        Limit {
+            stream,
+            remaining: amount,
+        }
+    }
+}
+
+impl<T> BufStream for Limit<T>
+where
+    T: BufStream,
+{
+    type Item = T::Item;
+    type Error = LimitError<T::Error>;
+
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        use futures::Async::Ready;
+
+        if self.stream.size_hint().lower() > self.remaining {
+            return Err(LimitError { inner: None });
+        }
+
+        let res = self
+            .stream
+            .poll_buf()
+            .map_err(|err| LimitError { inner: Some(err) });
+
+        match res {
+            Ok(Ready(Some(ref buf))) => {
+                if buf.remaining() as u64 > self.remaining {
+                    self.remaining = 0;
+                    return Err(LimitError { inner: None });
+                }
+
+                self.remaining -= buf.remaining() as u64;
+            }
+            _ => {}
+        }
+
+        res
+    }
+}
+
+// ===== impl LimitError =====
+
+impl<T> LimitError<T> {
+    /// Returns `true` if the error was caused by polling the stream.
+    pub fn is_stream_err(&self) -> bool {
+        self.inner.is_some()
+    }
+
+    /// Returns `true` if the stream reached its limit.
+    pub fn is_limit_err(&self) -> bool {
+        self.inner.is_none()
+    }
+}

tokio-buf/src/util/mod.rs

@@ -0,0 +1,73 @@
+//! Types and utilities for working with `BufStream`.
+
+mod chain;
+mod collect;
+mod from;
+mod limit;
+
+pub use self::chain::Chain;
+pub use self::collect::Collect;
+pub use self::from::FromBufStream;
+pub use self::limit::Limit;
+
+pub mod error {
+    //! Error types
+
+    pub use super::collect::CollectError;
+    pub use super::from::CollectVecError;
+    pub use super::limit::LimitError;
+}
+
+use BufStream;
+
+impl<T> BufStreamExt for T where T: BufStream {}
+
+/// An extension trait for `BufStream`'s that provides a variety of convenient
+/// adapters.
+pub trait BufStreamExt: BufStream {
+    /// Takes two buf streams and creates a new buf stream over both in
+    /// sequence.
+    ///
+    /// `chain()` returns a new `BufStream` value which will first yield all
+    /// data from `self` then all data from `other`.
+    ///
+    /// In other words, it links two buf streams together, in a chain.
+    fn chain<T>(self, other: T) -> Chain<Self, T>
+    where
+        Self: Sized,
+        T: BufStream<Error = Self::Error>,
+    {
+        Chain::new(self, other)
+    }
+
+    /// Consumes all data from `self`, storing it in byte storage of type `T`.
+    ///
+    /// `collect()` returns a future that buffers all data yielded from `self`
+    /// into storage of type of `T`. The future completes once `self` yield
+    /// `None`, returning the buffered data.
+    ///
+    /// The collect future will yield an error if `self` yields an error or if
+    /// the collect operation errors. The collect error cases are dependent on
+    /// the target storage type.
+    fn collect<T>(self) -> Collect<Self, T>
+    where
+        Self: Sized,
+        T: FromBufStream<Self::Item>,
+    {
+        Collect::new(self)
+    }
+
+    /// Limit the number of bytes that the stream can yield.
+    ///
+    /// `limit()` returns a new `BufStream` value which yields all the data from
+    /// `self` while ensuring that at most `amount` bytes are yielded.
+    ///
+    /// If `self` can yield greater than `amount` bytes, the returned stream
+    /// will yield an error.
+    fn limit(self, amount: u64) -> Limit<Self>
+    where
+        Self: Sized,
+    {
+        Limit::new(self, amount)
+    }
+}