tokio/io/util/
chain.rs

1use crate::io::{AsyncBufRead, AsyncRead, ReadBuf};
2
3use pin_project_lite::pin_project;
4use std::fmt;
5use std::io;
6use std::pin::Pin;
7use std::task::{ready, Context, Poll};
8
9pin_project! {
10    /// Stream for the [`chain`](super::AsyncReadExt::chain) method.
11    #[must_use = "streams do nothing unless polled"]
12    #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))]
13    pub struct Chain<T, U> {
14        #[pin]
15        first: T,
16        #[pin]
17        second: U,
18        done_first: bool,
19    }
20}
21
22pub(super) fn chain<T, U>(first: T, second: U) -> Chain<T, U>
23where
24    T: AsyncRead,
25    U: AsyncRead,
26{
27    Chain {
28        first,
29        second,
30        done_first: false,
31    }
32}
33
34impl<T, U> Chain<T, U>
35where
36    T: AsyncRead,
37    U: AsyncRead,
38{
39    /// Gets references to the underlying readers in this `Chain`.
40    pub fn get_ref(&self) -> (&T, &U) {
41        (&self.first, &self.second)
42    }
43
44    /// Gets mutable references to the underlying readers in this `Chain`.
45    ///
46    /// Care should be taken to avoid modifying the internal I/O state of the
47    /// underlying readers as doing so may corrupt the internal state of this
48    /// `Chain`.
49    pub fn get_mut(&mut self) -> (&mut T, &mut U) {
50        (&mut self.first, &mut self.second)
51    }
52
53    /// Gets pinned mutable references to the underlying readers in this `Chain`.
54    ///
55    /// Care should be taken to avoid modifying the internal I/O state of the
56    /// underlying readers as doing so may corrupt the internal state of this
57    /// `Chain`.
58    pub fn get_pin_mut(self: Pin<&mut Self>) -> (Pin<&mut T>, Pin<&mut U>) {
59        let me = self.project();
60        (me.first, me.second)
61    }
62
63    /// Consumes the `Chain`, returning the wrapped readers.
64    pub fn into_inner(self) -> (T, U) {
65        (self.first, self.second)
66    }
67}
68
69impl<T, U> fmt::Debug for Chain<T, U>
70where
71    T: fmt::Debug,
72    U: fmt::Debug,
73{
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        f.debug_struct("Chain")
76            .field("t", &self.first)
77            .field("u", &self.second)
78            .finish()
79    }
80}
81
82impl<T, U> AsyncRead for Chain<T, U>
83where
84    T: AsyncRead,
85    U: AsyncRead,
86{
87    fn poll_read(
88        self: Pin<&mut Self>,
89        cx: &mut Context<'_>,
90        buf: &mut ReadBuf<'_>,
91    ) -> Poll<io::Result<()>> {
92        let me = self.project();
93
94        if !*me.done_first {
95            let rem = buf.remaining();
96            ready!(me.first.poll_read(cx, buf))?;
97            // A read that fills nothing only indicates EOF if the buffer
98            // had capacity to begin with.
99            if buf.remaining() == rem && rem != 0 {
100                *me.done_first = true;
101            } else {
102                return Poll::Ready(Ok(()));
103            }
104        }
105        me.second.poll_read(cx, buf)
106    }
107}
108
109impl<T, U> AsyncBufRead for Chain<T, U>
110where
111    T: AsyncBufRead,
112    U: AsyncBufRead,
113{
114    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
115        let me = self.project();
116
117        if !*me.done_first {
118            match ready!(me.first.poll_fill_buf(cx)?) {
119                [] => {
120                    *me.done_first = true;
121                }
122                buf => return Poll::Ready(Ok(buf)),
123            }
124        }
125        me.second.poll_fill_buf(cx)
126    }
127
128    fn consume(self: Pin<&mut Self>, amt: usize) {
129        let me = self.project();
130        if !*me.done_first {
131            me.first.consume(amt)
132        } else {
133            me.second.consume(amt)
134        }
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn assert_unpin() {
144        crate::is_unpin::<Chain<(), ()>>();
145    }
146}
tokio/io/util/chain.rs

tokio/io/util/
chain.rs
