Add tokio-buf and a BufStream trait (#611)

The `BufStream` trait provides an improved API for working with
asynchronous streams of bytes compared to `Stream<Item = [u8]>`

tokio-buf/src/buf_stream/bytes.rs

@@ -0,0 +1,70 @@
+use BufStream;
+use buf_stream::errors::internal::Never;
+
+use bytes::{Bytes, BytesMut};
+use futures::Poll;
+
+use std::io;
+
+impl BufStream for Vec<u8> {
+    type Item = io::Cursor<Vec<u8>>;
+    type Error = Never;
+
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        if self.is_empty() {
+            return Ok(None.into());
+        }
+
+        poll_bytes(self)
+    }
+}
+
+impl BufStream for &'static [u8] {
+    type Item = io::Cursor<&'static [u8]>;
+    type Error = Never;
+
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        if self.is_empty() {
+            return Ok(None.into());
+        }
+
+        poll_bytes(self)
+    }
+}
+
+impl BufStream for Bytes {
+    type Item = io::Cursor<Bytes>;
+    type Error = Never;
+
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        if self.is_empty() {
+            return Ok(None.into());
+        }
+
+        poll_bytes(self)
+    }
+}
+
+impl BufStream for BytesMut {
+    type Item = io::Cursor<BytesMut>;
+    type Error = Never;
+
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        if self.is_empty() {
+            return Ok(None.into());
+        }
+
+        poll_bytes(self)
+    }
+}
+
+fn poll_bytes<T: Default>(buf: &mut T)
+    -> Poll<Option<io::Cursor<T>>, Never>
+{
+    use std::mem;
+
+    let bytes = mem::replace(buf, Default::default());
+    let buf = io::Cursor::new(bytes);
+
+    Ok(Some(buf).into())
+}

tokio-buf/src/buf_stream/chain.rs

@@ -0,0 +1,51 @@
+use super::{BufStream, SizeHint};
+
+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())
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        // TODO: Implement
+        SizeHint::default()
+    }
+}

tokio-buf/src/buf_stream/collect.rs

@@ -0,0 +1,103 @@
+use super::{BufStream, FromBufStream};
+
+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/buf_stream/errors.rs

@@ -0,0 +1,32 @@
+//! Error types
+
+pub use super::collect::CollectError;
+pub use super::from::CollectVecError;
+pub use super::limit::LimitError;
+
+// Being crate-private, we should be able to swap the type out in a
+// backwards compatible way.
+pub(crate) mod internal {
+    use std::{error, fmt};
+
+    /// An error that can never occur
+    pub enum Never {}
+
+    impl fmt::Debug for Never {
+        fn fmt(&self, _f: &mut fmt::Formatter) -> fmt::Result {
+            match *self {}
+        }
+    }
+
+    impl fmt::Display for Never {
+        fn fmt(&self, _f: &mut fmt::Formatter) -> fmt::Result {
+            match *self {}
+        }
+    }
+
+    impl error::Error for Never {
+        fn description(&self) -> &str {
+            match *self {}
+        }
+    }
+}

tokio-buf/src/buf_stream/from.rs

@@ -0,0 +1,110 @@
+use super::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/buf_stream/limit.rs

@@ -0,0 +1,92 @@
+use super::{BufStream, SizeHint};
+
+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
+    }
+
+    fn size_hint(&self) -> SizeHint {
+        let mut hint = self.stream.size_hint();
+
+        let upper = hint.upper()
+            .map(|upper| upper.min(self.remaining))
+            .unwrap_or(self.remaining);
+
+        hint.set_upper(upper);
+        hint
+    }
+
+    fn consume_hint(&mut self, amount: usize) {
+        // TODO: Should this be capped by `self.remaining`?
+        self.stream.consume_hint(amount)
+    }
+}
+
+// ===== 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/buf_stream/mod.rs

@@ -0,0 +1,163 @@
+//! Types and utilities for working with `BufStream`.
+
+mod bytes;
+mod chain;
+mod collect;
+pub mod errors;
+mod from;
+mod limit;
+mod size_hint;
+mod str;
+
+pub use self::chain::Chain;
+pub use self::collect::Collect;
+pub use self::from::FromBufStream;
+pub use self::limit::Limit;
+pub use self::size_hint::SizeHint;
+
+use bytes::Buf;
+use futures::Poll;
+
+/// An asynchronous stream of bytes.
+///
+/// `BufStream` asynchronously yields values implementing `Buf`, i.e. byte
+/// buffers.
+pub trait BufStream {
+    /// Values yielded by the `BufStream`.
+    ///
+    /// Each item is a sequence of bytes representing a chunk of the total
+    /// `ByteStream`.
+    type Item: Buf;
+
+    /// The error type this `BufStream` might generate.
+    type Error;
+
+    /// Attempt to pull out the next buffer of this stream, registering the
+    /// current task for wakeup if the value is not yet available, and returning
+    /// `None` if the stream is exhausted.
+    ///
+    /// # Return value
+    ///
+    /// There are several possible return values, each indicating a distinct
+    /// stream state:
+    ///
+    /// - `Ok(Async::NotReady)` means that this stream's next value is not ready
+    /// yet. Implementations will ensure that the current task will be notified
+    /// when the next value may be ready.
+    ///
+    /// - `Ok(Async::Ready(Some(buf)))` means that the stream has successfully
+    /// produced a value, `buf`, and may produce further values on subsequent
+    /// `poll_buf` calls.
+    ///
+    /// - `Ok(Async::Ready(None))` means that the stream has terminated, and
+    /// `poll_buf` should not be invoked again.
+    ///
+    /// # Panics
+    ///
+    /// Once a stream is finished, i.e. `Ready(None)` has been returned, further
+    /// calls to `poll_buf` may result in a panic or other "bad behavior".
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error>;
+
+    /// Returns the bounds on the remaining length of the stream.
+    ///
+    /// The size hint allows the caller to perform certain optimizations that
+    /// are dependent on the byte stream size. For example, `collect` uses the
+    /// size hint to pre-allocate enough capacity to store the entirety of the
+    /// data received from the byte stream.
+    ///
+    /// When `SizeHint::upper()` returns `Some` with a value equal to
+    /// `SizeHint::lower()`, this represents the exact number of bytes that will
+    /// be yielded by the `BufStream`.
+    ///
+    /// # Implementation notes
+    ///
+    /// While not enforced, implementations are expected to respect the values
+    /// returned from `SizeHint`. Any deviation is considered an implementation
+    /// bug. Consumers may rely on correctness in order to use the value as part
+    /// of protocol impelmentations. For example, an HTTP library may use the
+    /// size hint to set the `content-length` header.
+    ///
+    /// However, `size_hint` must not be trusted to omit bounds checks in unsafe
+    /// code. An incorrect implementation of `size_hint()` must not lead to
+    /// memory safety violations.
+    fn size_hint(&self) -> SizeHint {
+        SizeHint::default()
+    }
+
+    /// Indicates to the `BufStream` how much data the consumer is currently
+    /// able to process.
+    ///
+    /// The consume hint allows the stream to perform certain optimizations that
+    /// are dependent on the consumer's readiness. For example, the consume hint
+    /// may be used to request a remote peer to start sending up to `amount`
+    /// data.
+    ///
+    /// Calling `consume_hint` is not a requirement. If `consume_hint` is never
+    /// called, the stream should assume a default behavior. When `consume_hint`
+    /// is called, the stream should make a best effort to honor by the request.
+    ///
+    /// `amount` represents the number of bytes that the caller would like to
+    /// receive at the time the function is called. For example, if
+    /// `consume_hint` is called with 20, the consumer requests 20 bytes. The
+    /// stream may yield less than that. If the next call to `poll_buf` returns
+    /// 5 bytes, the consumer still has 15 bytes requested. At this point,
+    /// invoking `consume_hint` again with 20 resets the amount requested back
+    /// to 20 bytes.
+    ///
+    /// Calling `consume_hint` with 0 as the argument informs the stream that
+    /// the caller does not intend to call `poll_buf`. If `poll_buf` **is**
+    /// called, the stream may, but is not obligated to, return `NotReady` even
+    /// if it could produce data at that point. If it chooses to return
+    /// `NotReady`, when `consume_hint` is called with a non-zero argument, the
+    /// task must be notified in order to respect the `poll_buf` contract.
+    fn consume_hint(&mut self, amount: usize) {
+        // By default, this function does nothing
+        drop(amount);
+    }
+
+    /// 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)
+    }
+}

tokio-buf/src/buf_stream/size_hint.rs

@@ -0,0 +1,56 @@
+use std::u64;
+
+/// A `BufStream` size hint
+///
+/// The default implementation returns:
+///
+/// * 0 for `available`
+/// * 0 for `lower`
+/// * `None` for `upper`.
+#[derive(Debug, Default, Clone)]
+pub struct SizeHint {
+    lower: u64,
+    upper: Option<u64>,
+}
+
+impl SizeHint {
+    /// Returns a new `SizeHint` with default values
+    pub fn new() -> SizeHint {
+        SizeHint::default()
+    }
+
+    /// Returns the lower bound of data that the `BufStream` will yield before
+    /// completing.
+    pub fn lower(&self) -> u64 {
+        self.lower
+    }
+
+    /// Set the value of the `lower` hint.
+    ///
+    /// # Panics
+    ///
+    /// The function panics if `value` is less than `upper`.
+    pub fn set_lower(&mut self, value: u64) {
+        assert!(value <= self.upper.unwrap_or(u64::MAX));
+        self.lower = value;
+    }
+
+    /// Returns the upper bound of data the `BufStream` will yield before
+    /// completing, or `None` if the value is unknown.
+    pub fn upper(&self) -> Option<u64> {
+        self.upper
+    }
+
+    /// Set the value of the `upper` hint value.
+    ///
+    /// # Panics
+    ///
+    /// This function panics if `value` is less than `lower`.
+    pub fn set_upper(&mut self, value: u64) {
+        // There is no need to check `available` as that is guaranteed to be
+        // less than or equal to `lower`.
+        assert!(value >= self.lower, "`value` is less than than `lower`");
+
+        self.upper = Some(value);
+    }
+}

tokio-buf/src/buf_stream/str.rs

@@ -0,0 +1,39 @@
+use BufStream;
+use buf_stream::errors::internal::Never;
+
+use futures::Poll;
+
+use std::io;
+use std::mem;
+
+impl BufStream for String {
+    type Item = io::Cursor<Vec<u8>>;
+    type Error = Never;
+
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        if self.is_empty() {
+            return Ok(None.into());
+        }
+
+        let bytes = mem::replace(self, Default::default()).into_bytes();
+        let buf = io::Cursor::new(bytes);
+
+        Ok(Some(buf).into())
+    }
+}
+
+impl BufStream for &'static str {
+    type Item = io::Cursor<&'static [u8]>;
+    type Error = Never;
+
+    fn poll_buf(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
+        if self.is_empty() {
+            return Ok(None.into());
+        }
+
+        let bytes = mem::replace(self, Default::default()).as_bytes();
+        let buf = io::Cursor::new(bytes);
+
+        Ok(Some(buf).into())
+    }
+}

tokio-buf/src/lib.rs

@@ -0,0 +1,20 @@
+#![doc(html_root_url = "https://docs.rs/tokio-buf/0.1.0")]
+#![deny(missing_docs, missing_debug_implementations)]
+#![cfg_attr(test, deny(warnings))]
+
+//! Asynchronous stream of bytes.
+//!
+//! This crate contains the `BufStream` trait and a number of combinators for
+//! this trait. The trait is similar to `Stream` in the `futures` library, but
+//! instead of yielding arbitrary values, it only yields types that implement
+//! `Buf` (i.e, byte collections).
+
+extern crate bytes;
+extern crate either;
+#[macro_use]
+extern crate futures;
+
+pub mod buf_stream;
+
+#[doc(inline)]
+pub use buf_stream::BufStream;