compio_io\util/repeat.rs
1use std::mem::MaybeUninit;
2
3use compio_buf::{BufResult, IoVectoredBufMut, SetLenExt};
4
5use crate::{AsyncBufRead, AsyncRead, IoResult};
6
7/// A reader that infinitely repeats one byte constructed via [`repeat`].
8///
9/// All reads from this reader will succeed by filling the specified buffer with
10/// the given byte.
11///
12/// # Examples
13///
14/// ```rust
15/// # futures_executor::block_on(async {
16/// use compio_io::{self, AsyncRead, AsyncReadExt};
17///
18/// let (len, buffer) = compio_io::repeat(42)
19/// .read(Vec::with_capacity(3))
20/// .await
21/// .unwrap();
22///
23/// assert_eq!(buffer.as_slice(), [42, 42, 42]);
24/// assert_eq!(len, 3);
25/// # })
26/// ```
27pub struct Repeat(u8);
28
29impl AsyncRead for Repeat {
30 async fn read<B: compio_buf::IoBufMut>(
31 &mut self,
32 mut buf: B,
33 ) -> compio_buf::BufResult<usize, B> {
34 let slice = buf.as_uninit();
35
36 let len = slice.len();
37 slice.fill(MaybeUninit::new(self.0));
38 // SAFETY: we just initialized exactly `len` bytes in `buf` from index
39 // 0.
40 unsafe { buf.advance(len) };
41
42 BufResult(Ok(len), buf)
43 }
44
45 async fn read_vectored<V: IoVectoredBufMut>(&mut self, mut buf: V) -> BufResult<usize, V> {
46 let mut len: usize = 0;
47 for slice in buf.iter_uninit_slice() {
48 len = len
49 .checked_add(slice.len())
50 .expect("total vectored buffer length overflow");
51 slice.fill(MaybeUninit::new(self.0));
52 }
53 debug_assert_eq!(len, buf.total_capacity());
54 // SAFETY: every byte counted in `len` is initialized in the loop above.
55 unsafe { buf.advance_vec_to(len) };
56
57 BufResult(Ok(len), buf)
58 }
59}
60
61impl AsyncBufRead for Repeat {
62 async fn fill_buf(&mut self) -> IoResult<&'_ [u8]> {
63 Ok(std::slice::from_ref(&self.0))
64 }
65
66 fn consume(&mut self, _: usize) {}
67}
68
69/// Creates a reader that infinitely repeats one byte.
70///
71/// All reads from this reader will succeed by filling the specified buffer with
72/// the given byte.
73///
74/// # Examples
75///
76/// ```rust
77/// # futures_executor::block_on(async {
78/// use compio_io::{self, AsyncRead, AsyncReadExt};
79///
80/// let ((), buffer) = compio_io::repeat(42)
81/// .read_exact(Vec::with_capacity(3))
82/// .await
83/// .unwrap();
84///
85/// assert_eq!(buffer.as_slice(), [42, 42, 42]);
86/// # })
87/// ```
88pub fn repeat(byte: u8) -> Repeat {
89 Repeat(byte)
90}