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