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}