madsim_real_tokio/io/util/
repeat.rs

1use crate::io::util::poll_proceed_and_make_progress;
2use crate::io::{AsyncRead, ReadBuf};
3
4use std::io;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8cfg_io_util! {
9    /// An async reader which yields one byte over and over and over and over and
10    /// over and...
11    ///
12    /// This struct is generally created by calling [`repeat`][repeat]. Please
13    /// see the documentation of `repeat()` for more details.
14    ///
15    /// This is an asynchronous version of [`std::io::Repeat`][std].
16    ///
17    /// [repeat]: fn@repeat
18    /// [std]: std::io::Repeat
19    #[derive(Debug)]
20    pub struct Repeat {
21        byte: u8,
22    }
23
24    /// Creates an instance of an async reader that infinitely repeats one byte.
25    ///
26    /// All reads from this reader will succeed by filling the specified buffer with
27    /// the given byte.
28    ///
29    /// This is an asynchronous version of [`std::io::repeat`][std].
30    ///
31    /// [std]: std::io::repeat
32    ///
33    /// # Examples
34    ///
35    /// ```
36    /// use tokio::io::{self, AsyncReadExt};
37    ///
38    /// #[tokio::main]
39    /// async fn main() {
40    ///     let mut buffer = [0; 3];
41    ///     io::repeat(0b101).read_exact(&mut buffer).await.unwrap();
42    ///     assert_eq!(buffer, [0b101, 0b101, 0b101]);
43    /// }
44    /// ```
45    pub fn repeat(byte: u8) -> Repeat {
46        Repeat { byte }
47    }
48}
49
50impl AsyncRead for Repeat {
51    #[inline]
52    fn poll_read(
53        self: Pin<&mut Self>,
54        cx: &mut Context<'_>,
55        buf: &mut ReadBuf<'_>,
56    ) -> Poll<io::Result<()>> {
57        ready!(crate::trace::trace_leaf(cx));
58        ready!(poll_proceed_and_make_progress(cx));
59        // TODO: could be faster, but should we unsafe it?
60        while buf.remaining() != 0 {
61            buf.put_slice(&[self.byte]);
62        }
63        Poll::Ready(Ok(()))
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn assert_unpin() {
73        crate::is_unpin::<Repeat>();
74    }
75}