      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
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121

    use crate::io::util::read_until::read_until_internal;
    use crate::io::AsyncBufRead;

    use pin_project_lite::pin_project;
    use std::io;
    use std::mem;
    use std::pin::Pin;
    use std::task::{Context, Poll};

    pin_project! {
        /// Splitter for the [`split`](crate::io::AsyncBufReadExt::split) method.
        ///
        /// A `Split` can be turned into a `Stream` with [`SplitStream`].
        ///
        /// [`SplitStream`]: https://docs.rs/tokio-stream/0.1/tokio_stream/wrappers/struct.SplitStream.html
        #[derive(Debug)]
        #[must_use = "streams do nothing unless polled"]
        #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
        pub struct Split<R> {
            #[pin]
            reader: R,
            buf: Vec<u8>,
            delim: u8,
            read: usize,
        }
    }

    pub(crate) fn split<R>(reader: R, delim: u8) -> Split<R>
    where
        R: AsyncBufRead,
    {
        Split {
            reader,
            buf: Vec::new(),
            delim,
            read: 0,
        }
    }

    impl<R> Split<R>
    where
        R: AsyncBufRead + Unpin,
    {
        /// Returns the next segment in the stream.
        ///
        /// # Examples
        ///
        /// ```
        /// # use tokio::io::AsyncBufRead;
        /// use tokio::io::AsyncBufReadExt;
        ///
        /// # async fn dox(my_buf_read: impl AsyncBufRead + Unpin) -> std::io::Result<()> {
        /// let mut segments = my_buf_read.split(b'f');
        ///
        /// while let Some(segment) = segments.next_segment().await? {
        ///     println!("length = {}", segment.len())
        /// }
        /// # Ok(())
        /// # }
        /// ```
        pub async fn next_segment(&mut self) -> io::Result<Option<Vec<u8>>> {
            use crate::future::poll_fn;

            poll_fn(|cx| Pin::new(&mut *self).poll_next_segment(cx)).await
        }
    }

    impl<R> Split<R>
    where
        R: AsyncBufRead,
    {
        /// Polls for the next segment in the stream.
        ///
        /// This method returns:
        ///
        ///  * `Poll::Pending` if the next segment is not yet available.
        ///  * `Poll::Ready(Ok(Some(segment)))` if the next segment is available.
        ///  * `Poll::Ready(Ok(None))` if there are no more segments in this stream.
        ///  * `Poll::Ready(Err(err))` if an IO error occurred while reading the
        ///    next segment.
        ///
        /// When the method returns `Poll::Pending`, the `Waker` in the provided
        /// `Context` is scheduled to receive a wakeup when more bytes become
        /// available on the underlying IO resource.
        ///
        /// Note that on multiple calls to `poll_next_segment`, only the `Waker`
        /// from the `Context` passed to the most recent call is scheduled to
        /// receive a wakeup.
        pub fn poll_next_segment(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
        ) -> Poll<io::Result<Option<Vec<u8>>>> {
            let me = self.project();

            let n = ready!(read_until_internal(
                me.reader, cx, *me.delim, me.buf, me.read,
            ))?;
            // read_until_internal resets me.read to zero once it finds the delimeter
            debug_assert_eq!(*me.read, 0);

            if n == 0 && me.buf.is_empty() {
                return Poll::Ready(Ok(None));
            }

            if me.buf.last() == Some(me.delim) {
                me.buf.pop();
            }

            Poll::Ready(Ok(Some(mem::replace(me.buf, Vec::new()))))
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;

        #[test]
        fn assert_unpin() {
            crate::is_unpin::<Split<()>>();
        }
    }