1 use crate::io::util::read_until::read_until_internal;
2 use crate::io::AsyncBufRead;
3 
4 use pin_project_lite::pin_project;
5 use std::io;
6 use std::mem;
7 use std::pin::Pin;
8 use std::task::{Context, Poll};
9 
10 pin_project! {
11     /// Stream for the [`split`](crate::io::AsyncBufReadExt::split) method.
12     #[derive(Debug)]
13     #[must_use = "streams do nothing unless polled"]
14     #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
15     pub struct Split<R> {
16         #[pin]
17         reader: R,
18         buf: Vec<u8>,
19         delim: u8,
20         read: usize,
21     }
22 }
23 
split<R>(reader: R, delim: u8) -> Split<R> where R: AsyncBufRead,24 pub(crate) fn split<R>(reader: R, delim: u8) -> Split<R>
25 where
26     R: AsyncBufRead,
27 {
28     Split {
29         reader,
30         buf: Vec::new(),
31         delim,
32         read: 0,
33     }
34 }
35 
36 impl<R> Split<R>
37 where
38     R: AsyncBufRead + Unpin,
39 {
40     /// Returns the next segment in the stream.
41     ///
42     /// # Examples
43     ///
44     /// ```
45     /// # use tokio::io::AsyncBufRead;
46     /// use tokio::io::AsyncBufReadExt;
47     ///
48     /// # async fn dox(my_buf_read: impl AsyncBufRead + Unpin) -> std::io::Result<()> {
49     /// let mut segments = my_buf_read.split(b'f');
50     ///
51     /// while let Some(segment) = segments.next_segment().await? {
52     ///     println!("length = {}", segment.len())
53     /// }
54     /// # Ok(())
55     /// # }
56     /// ```
next_segment(&mut self) -> io::Result<Option<Vec<u8>>>57     pub async fn next_segment(&mut self) -> io::Result<Option<Vec<u8>>> {
58         use crate::future::poll_fn;
59 
60         poll_fn(|cx| Pin::new(&mut *self).poll_next_segment(cx)).await
61     }
62 }
63 
64 impl<R> Split<R>
65 where
66     R: AsyncBufRead,
67 {
68     #[doc(hidden)]
poll_next_segment( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll<io::Result<Option<Vec<u8>>>>69     pub fn poll_next_segment(
70         self: Pin<&mut Self>,
71         cx: &mut Context<'_>,
72     ) -> Poll<io::Result<Option<Vec<u8>>>> {
73         let me = self.project();
74 
75         let n = ready!(read_until_internal(
76             me.reader, cx, *me.delim, me.buf, me.read,
77         ))?;
78         // read_until_internal resets me.read to zero once it finds the delimeter
79         debug_assert_eq!(*me.read, 0);
80 
81         if n == 0 && me.buf.is_empty() {
82             return Poll::Ready(Ok(None));
83         }
84 
85         if me.buf.last() == Some(me.delim) {
86             me.buf.pop();
87         }
88 
89         Poll::Ready(Ok(Some(mem::replace(me.buf, Vec::new()))))
90     }
91 }
92 
93 #[cfg(feature = "stream")]
94 impl<R: AsyncBufRead> crate::stream::Stream for Split<R> {
95     type Item = io::Result<Vec<u8>>;
96 
poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>>97     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
98         Poll::Ready(match ready!(self.poll_next_segment(cx)) {
99             Ok(Some(segment)) => Some(Ok(segment)),
100             Ok(None) => None,
101             Err(err) => Some(Err(err)),
102         })
103     }
104 }
105 
106 #[cfg(test)]
107 mod tests {
108     use super::*;
109 
110     #[test]
assert_unpin()111     fn assert_unpin() {
112         crate::is_unpin::<Split<()>>();
113     }
114 }
115