1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
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,
        }
    }
}