1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
use std::mem::MaybeUninit;

use compio_buf::BufResult;

use crate::{AsyncBufRead, AsyncRead, IoResult};

/// A reader that infinitely repeats one byte constructed via [`repeat`].
///
/// All reads from this reader will succeed by filling the specified buffer with
/// the given byte.
///
/// # Examples
///
/// ```rust
/// # compio_runtime::Runtime::new().unwrap().block_on(async {
/// use compio_io::{self, AsyncRead, AsyncReadExt};
///
/// let (len, buffer) = compio_io::repeat(42)
///     .read(Vec::with_capacity(3))
///     .await
///     .unwrap();
///
/// assert_eq!(buffer.as_slice(), [42, 42, 42]);
/// assert_eq!(len, 3);
/// # })
/// ```
pub struct Repeat(u8);

impl AsyncRead for Repeat {
    async fn read<B: compio_buf::IoBufMut>(
        &mut self,
        mut buf: B,
    ) -> compio_buf::BufResult<usize, B> {
        let slice = buf.as_mut_slice();

        let len = slice.len();
        slice.fill(MaybeUninit::new(self.0));
        unsafe { buf.set_buf_init(len) };

        BufResult(Ok(len), buf)
    }
}

impl AsyncBufRead for Repeat {
    async fn fill_buf(&mut self) -> IoResult<&'_ [u8]> {
        Ok(std::slice::from_ref(&self.0))
    }

    fn consume(&mut self, _: usize) {}
}

/// Creates a reader that infinitely repeats one byte.
///
/// All reads from this reader will succeed by filling the specified buffer with
/// the given byte.
///
/// # Examples
///
/// ```rust
/// # compio_runtime::Runtime::new().unwrap().block_on(async {
/// use compio_io::{self, AsyncRead, AsyncReadExt};
///
/// let ((), buffer) = compio_io::repeat(42)
///     .read_exact(Vec::with_capacity(3))
///     .await
///     .unwrap();
///
/// assert_eq!(buffer.as_slice(), [42, 42, 42]);
/// # })
/// ```
pub fn repeat(byte: u8) -> Repeat {
    Repeat(byte)
}