Skip to main content

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        // SAFETY: The whole slice is filled with initialized values below.
35        let slice = unsafe { buf.as_uninit() };
36
37        let len = slice.len();
38        slice.fill(MaybeUninit::new(self.0));
39        // SAFETY: we just initialized exactly `len` bytes in `buf` from index
40        // 0.
41        unsafe { buf.advance(len) };
42
43        BufResult(Ok(len), buf)
44    }
45
46    async fn read_vectored<V: IoVectoredBufMut>(&mut self, mut buf: V) -> BufResult<usize, V> {
47        let mut len: usize = 0;
48        // SAFETY: The loop fills every slice with initialized values.
49        for slice in unsafe { buf.iter_uninit_slice() } {
50            len = len
51                .checked_add(slice.len())
52                .expect("total vectored buffer length overflow");
53            slice.fill(MaybeUninit::new(self.0));
54        }
55        debug_assert_eq!(len, buf.total_capacity());
56        // SAFETY: every byte counted in `len` is initialized in the loop above.
57        unsafe { buf.advance_vec_to(len) };
58
59        BufResult(Ok(len), buf)
60    }
61}
62
63impl AsyncBufRead for Repeat {
64    async fn fill_buf(&mut self) -> IoResult<&'_ [u8]> {
65        Ok(std::slice::from_ref(&self.0))
66    }
67
68    fn consume(&mut self, _: usize) {}
69}
70
71/// Creates a reader that infinitely repeats one byte.
72///
73/// All reads from this reader will succeed by filling the specified buffer with
74/// the given byte.
75///
76/// # Examples
77///
78/// ```rust
79/// # futures_executor::block_on(async {
80/// use compio_io::{self, AsyncRead, AsyncReadExt};
81///
82/// let ((), buffer) = compio_io::repeat(42)
83///     .read_exact(Vec::with_capacity(3))
84///     .await
85///     .unwrap();
86///
87/// assert_eq!(buffer.as_slice(), [42, 42, 42]);
88/// # })
89/// ```
90pub fn repeat(byte: u8) -> Repeat {
91    Repeat(byte)
92}