broker_tokio/io/util/
split.rs

1use crate::io::util::read_until::read_until_internal;
2use crate::io::AsyncBufRead;
3
4use pin_project_lite::pin_project;
5use std::io;
6use std::mem;
7use std::pin::Pin;
8use std::task::{Context, Poll};
9
10pin_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
24pub(crate) fn split<R>(reader: R, delim: u8) -> Split<R>
25where
26    R: AsyncBufRead,
27{
28    Split {
29        reader,
30        buf: Vec::new(),
31        delim,
32        read: 0,
33    }
34}
35
36impl<R> Split<R>
37where
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    /// ```
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
64impl<R> Split<R>
65where
66    R: AsyncBufRead,
67{
68    #[doc(hidden)]
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
79        if n == 0 && me.buf.is_empty() {
80            return Poll::Ready(Ok(None));
81        }
82
83        if me.buf.last() == Some(me.delim) {
84            me.buf.pop();
85        }
86
87        Poll::Ready(Ok(Some(mem::replace(me.buf, Vec::new()))))
88    }
89}
90
91#[cfg(feature = "stream")]
92impl<R: AsyncBufRead> crate::stream::Stream for Split<R> {
93    type Item = io::Result<Vec<u8>>;
94
95    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
96        Poll::Ready(match ready!(self.poll_next_segment(cx)) {
97            Ok(Some(segment)) => Some(Ok(segment)),
98            Ok(None) => None,
99            Err(err) => Some(Err(err)),
100        })
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn assert_unpin() {
110        crate::is_unpin::<Split<()>>();
111    }
112}
broker_tokio/io/util/split.rs

broker_tokio/io/util/
split.rs
