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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
#[cfg(feature = "allocator_api")]
use std::alloc::Allocator;
use std::{io::Cursor, ops::DerefMut, rc::Rc, sync::Arc};

use compio_buf::{buf_try, t_alloc, BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBufMut};

mod buf;
#[macro_use]
mod ext;

pub use buf::*;
pub use ext::*;

use crate::util::slice_to_buf;

/// AsyncRead
///
/// Async read with a ownership of a buffer
pub trait AsyncRead {
    /// Read some bytes from this source into the [`IoBufMut`] buffer and return
    /// a [`BufResult`], consisting of the buffer and a [`usize`] indicating
    /// how many bytes were read.
    ///
    /// # Caution
    ///
    /// Implementor **MUST** update the buffer init via
    /// [`SetBufInit::set_buf_init`] after reading, and no further update should
    /// be made by caller.
    ///
    /// [`SetBufInit::set_buf_init`]: compio_buf::SetBufInit::set_buf_init
    async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B>;

    /// Like `read`, except that it reads into a type implements
    /// [`IoVectoredBufMut`].
    ///
    /// The default implementation will try to read into the buffers in order,
    /// and stop whenever the reader returns an error, `Ok(0)`, or a length
    /// less than the length of the buf passed in, meaning it's possible that
    /// not all buffer space is filled. If guaranteed full read is desired,
    /// it is recommended to use [`AsyncReadExt::read_vectored_exact`]
    /// instead.
    ///
    /// # Caution
    ///
    /// Implementor **MUST** update the buffer init via
    /// [`SetBufInit::set_buf_init`] after reading.
    ///
    /// [`SetBufInit::set_buf_init`]: compio_buf::SetBufInit::set_buf_init
    async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
        loop_read_vectored!(
            buf, len, total: usize, n, iter,
            loop self.read(iter),
            break if n == 0 || n < len {
                Some(Ok(total))
            } else {
                None
            }
        )
    }
}

impl<A: AsyncRead + ?Sized> AsyncRead for &mut A {
    #[inline(always)]
    async fn read<T: IoBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
        (**self).read(buf).await
    }

    #[inline(always)]
    async fn read_vectored<T: IoVectoredBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
        (**self).read_vectored(buf).await
    }
}

impl<R: AsyncRead + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator> AsyncRead
    for t_alloc!(Box, R, A)
{
    #[inline(always)]
    async fn read<T: IoBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
        (**self).read(buf).await
    }

    #[inline(always)]
    async fn read_vectored<T: IoVectoredBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
        (**self).read_vectored(buf).await
    }
}

impl AsyncRead for &[u8] {
    #[inline]
    async fn read<T: IoBufMut>(&mut self, mut buf: T) -> BufResult<usize, T> {
        let len = slice_to_buf(self, &mut buf);
        *self = &self[len..];
        BufResult(Ok(len), buf)
    }

    async fn read_vectored<T: IoVectoredBufMut>(&mut self, mut buf: T) -> BufResult<usize, T> {
        let mut this = *self; // An immutable slice to track the read position

        for mut buf in buf.iter_buf_mut() {
            let n = slice_to_buf(this, buf.deref_mut());
            this = &this[n..];
            if this.is_empty() {
                break;
            }
        }

        BufResult(Ok(self.len() - this.len()), buf)
    }
}

/// # AsyncReadAt
///
/// Async read with a ownership of a buffer and a position
pub trait AsyncReadAt {
    /// Like [`AsyncRead::read`], except that it reads at a specified position.
    async fn read_at<T: IoBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T>;

    /// Like [`AsyncRead::read_vectored`], except that it reads at a specified
    /// position.
    async fn read_vectored_at<T: IoVectoredBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T> {
        loop_read_vectored!(
            buf, len, total: u64, n, iter,
            loop self.read_at(iter, pos + total),
            break if n == 0 || n < len {
                Some(Ok(total as usize))
            } else {
                None
            }
        )
    }
}

macro_rules! impl_read_at {
    (@ptr $($ty:ty),*) => {
        $(
            impl<A: AsyncReadAt + ?Sized> AsyncReadAt for $ty {
                async fn read_at<T: IoBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T> {
                    (**self).read_at(buf, pos).await
                }

                async fn read_vectored_at<T: IoVectoredBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T> {
                    (**self).read_vectored_at(buf, pos).await
                }
            }
        )*
    };

    (@ptra $($ty:ident),*) => {
        $(
            #[cfg(feature = "allocator_api")]
            impl<R: AsyncReadAt + ?Sized, A: Allocator> AsyncReadAt for $ty<R, A> {
                async fn read_at<T: IoBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T> {
                    (**self).read_at(buf, pos).await
                }

                async fn read_vectored_at<T: IoVectoredBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T> {
                    (**self).read_vectored_at(buf, pos).await
                }
            }
            #[cfg(not(feature = "allocator_api"))]
            impl_read_at!(@ptr $ty<A>);
        )*
    };

    (@slice $($(const $len:ident =>)? $ty:ty), *) => {
        $(
            impl<$(const $len: usize)?> AsyncReadAt for $ty {
                async fn read_at<T: IoBufMut>(&self, mut buf: T, pos: u64) -> BufResult<usize, T> {
                    let pos = pos.min(self.len() as u64);
                    let len = slice_to_buf(&self[pos as usize..], &mut buf);
                    BufResult(Ok(len), buf)
                }
            }
        )*
    }
}

impl_read_at!(@ptr &A, &mut A);
impl_read_at!(@ptra Box, Rc, Arc);
impl_read_at!(@slice [u8], const LEN => [u8; LEN]);

impl<#[cfg(feature = "allocator_api")] A: Allocator> AsyncReadAt for t_alloc!(Vec, u8, A) {
    async fn read_at<T: IoBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T> {
        self.as_slice().read_at(buf, pos).await
    }

    async fn read_vectored_at<T: IoVectoredBufMut>(&self, buf: T, pos: u64) -> BufResult<usize, T> {
        self.as_slice().read_vectored_at(buf, pos).await
    }
}

impl<A: AsyncReadAt> AsyncRead for Cursor<A> {
    #[inline]
    async fn read<T: IoBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
        let pos = self.position();
        let (n, buf) = buf_try!(self.get_ref().read_at(buf, pos).await);
        self.set_position(pos + n as u64);
        BufResult(Ok(n), buf)
    }

    #[inline]
    async fn read_vectored<T: IoVectoredBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
        let pos = self.position();
        let (n, buf) = buf_try!(self.get_ref().read_vectored_at(buf, pos).await);
        self.set_position(pos + n as u64);
        BufResult(Ok(n), buf)
    }
}