compio_io/read/
ext.rs

1#[cfg(feature = "allocator_api")]
2use std::alloc::Allocator;
3use std::{io, io::ErrorKind};
4
5use compio_buf::{BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBufMut, Uninit, t_alloc};
6
7use crate::{AsyncRead, AsyncReadAt, IoResult, util::Take};
8
9/// Shared code for read a scalar value from the underlying reader.
10macro_rules! read_scalar {
11    ($t:ty, $be:ident, $le:ident) => {
12        ::paste::paste! {
13            #[doc = concat!("Read a big endian `", stringify!($t), "` from the underlying reader.")]
14            async fn [< read_ $t >](&mut self) -> IoResult<$t> {
15                use ::compio_buf::{arrayvec::ArrayVec, BufResult};
16
17                const LEN: usize = ::std::mem::size_of::<$t>();
18                let BufResult(res, buf) = self.read_exact(ArrayVec::<u8, LEN>::new()).await;
19                res?;
20                // Safety: We just checked that the buffer is the correct size
21                Ok($t::$be(unsafe { buf.into_inner_unchecked() }))
22            }
23
24            #[doc = concat!("Read a little endian `", stringify!($t), "` from the underlying reader.")]
25            async fn [< read_ $t _le >](&mut self) -> IoResult<$t> {
26                use ::compio_buf::{arrayvec::ArrayVec, BufResult};
27
28                const LEN: usize = ::std::mem::size_of::<$t>();
29                let BufResult(res, buf) = self.read_exact(ArrayVec::<u8, LEN>::new()).await;
30                res?;
31                // Safety: We just checked that the buffer is the correct size
32                Ok($t::$le(unsafe { buf.into_inner_unchecked() }))
33            }
34        }
35    };
36}
37
38/// Shared code for loop reading until reaching a certain length.
39macro_rules! loop_read_exact {
40    ($buf:ident, $len:expr, $tracker:ident,loop $read_expr:expr) => {
41        let mut $tracker = 0;
42        let len = $len;
43
44        while $tracker < len {
45            match $read_expr.await.into_inner() {
46                BufResult(Ok(0), buf) => {
47                    return BufResult(
48                        Err(::std::io::Error::new(
49                            ::std::io::ErrorKind::UnexpectedEof,
50                            "failed to fill whole buffer",
51                        )),
52                        buf,
53                    );
54                }
55                BufResult(Ok(n), buf) => {
56                    $tracker += n;
57                    $buf = buf;
58                }
59                BufResult(Err(ref e), buf) if e.kind() == ::std::io::ErrorKind::Interrupted => {
60                    $buf = buf;
61                }
62                BufResult(Err(e), buf) => return BufResult(Err(e), buf),
63            }
64        }
65        return BufResult(Ok(()), $buf)
66    };
67}
68
69macro_rules! loop_read_vectored {
70    ($buf:ident, $iter:ident, $read_expr:expr) => {{
71        let mut $iter = match $buf.owned_iter() {
72            Ok(buf) => buf,
73            Err(buf) => return BufResult(Ok(0), buf),
74        };
75
76        loop {
77            let len = $iter.buf_capacity();
78            if len > 0 {
79                return $read_expr.await.into_inner();
80            }
81
82            match $iter.next() {
83                Ok(next) => $iter = next,
84                Err(buf) => return BufResult(Ok(0), buf),
85            }
86        }
87    }};
88}
89
90macro_rules! loop_read_to_end {
91    ($buf:ident, $tracker:ident : $tracker_ty:ty,loop $read_expr:expr) => {{
92        let mut $tracker: $tracker_ty = 0;
93        loop {
94            if $buf.len() == $buf.capacity() {
95                $buf.reserve(32);
96            }
97            match $read_expr.await.into_inner() {
98                BufResult(Ok(0), buf) => {
99                    $buf = buf;
100                    break;
101                }
102                BufResult(Ok(read), buf) => {
103                    $tracker += read as $tracker_ty;
104                    $buf = buf;
105                }
106                BufResult(Err(ref e), buf) if e.kind() == ::std::io::ErrorKind::Interrupted => {
107                    $buf = buf
108                }
109                res => return res,
110            }
111        }
112        BufResult(Ok($tracker as usize), $buf)
113    }};
114}
115
116#[inline]
117fn after_read_to_string(res: io::Result<usize>, buf: Vec<u8>) -> BufResult<usize, String> {
118    match res {
119        Err(err) => {
120            // we have to clear the read bytes if it is not valid utf8 bytes
121            let buf = String::from_utf8(buf).unwrap_or_else(|err| {
122                let mut buf = err.into_bytes();
123                buf.clear();
124
125                // Safety: the buffer is empty
126                unsafe { String::from_utf8_unchecked(buf) }
127            });
128
129            BufResult(Err(err), buf)
130        }
131        Ok(n) => match String::from_utf8(buf) {
132            Err(err) => BufResult(
133                Err(std::io::Error::new(ErrorKind::InvalidData, err)),
134                String::new(),
135            ),
136            Ok(data) => BufResult(Ok(n), data),
137        },
138    }
139}
140
141/// Implemented as an extension trait, adding utility methods to all
142/// [`AsyncRead`] types. Callers will tend to import this trait instead of
143/// [`AsyncRead`].
144pub trait AsyncReadExt: AsyncRead {
145    /// Creates a "by reference" adaptor for this instance of [`AsyncRead`].
146    ///
147    /// The returned adapter also implements [`AsyncRead`] and will simply
148    /// borrow this current reader.
149    fn by_ref(&mut self) -> &mut Self
150    where
151        Self: Sized,
152    {
153        self
154    }
155
156    /// Same as [`AsyncRead::read`], but it appends data to the end of the
157    /// buffer; in other words, it read to the beginning of the uninitialized
158    /// area.
159    async fn append<T: IoBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
160        self.read(buf.uninit()).await.map_buffer(Uninit::into_inner)
161    }
162
163    /// Read the exact number of bytes required to fill the buf.
164    async fn read_exact<T: IoBufMut>(&mut self, mut buf: T) -> BufResult<(), T> {
165        loop_read_exact!(buf, buf.buf_capacity(), read, loop self.read(buf.slice(read..)));
166    }
167
168    /// Read all bytes as [`String`] until underlying reader reaches `EOF`.
169    async fn read_to_string(&mut self, buf: String) -> BufResult<usize, String> {
170        let BufResult(res, buf) = self.read_to_end(buf.into_bytes()).await;
171        after_read_to_string(res, buf)
172    }
173
174    /// Read all bytes until underlying reader reaches `EOF`.
175    async fn read_to_end<#[cfg(feature = "allocator_api")] A: Allocator + 'static>(
176        &mut self,
177        mut buf: t_alloc!(Vec, u8, A),
178    ) -> BufResult<usize, t_alloc!(Vec, u8, A)> {
179        loop_read_to_end!(buf, total: usize, loop self.read(buf.slice(total..)))
180    }
181
182    /// Read the exact number of bytes required to fill the vectored buf.
183    async fn read_vectored_exact<T: IoVectoredBufMut>(&mut self, mut buf: T) -> BufResult<(), T> {
184        let len = buf.total_capacity();
185        loop_read_exact!(buf, len, read, loop self.read_vectored(buf.slice_mut(read)));
186    }
187
188    /// Creates an adaptor which reads at most `limit` bytes from it.
189    ///
190    /// This function returns a new instance of `AsyncRead` which will read
191    /// at most `limit` bytes, after which it will always return EOF
192    /// (`Ok(0)`). Any read errors will not count towards the number of
193    /// bytes read and future calls to [`read()`] may succeed.
194    ///
195    /// [`read()`]: AsyncRead::read
196    fn take(self, limit: u64) -> Take<Self>
197    where
198        Self: Sized,
199    {
200        Take::new(self, limit)
201    }
202
203    read_scalar!(u8, from_be_bytes, from_le_bytes);
204    read_scalar!(u16, from_be_bytes, from_le_bytes);
205    read_scalar!(u32, from_be_bytes, from_le_bytes);
206    read_scalar!(u64, from_be_bytes, from_le_bytes);
207    read_scalar!(u128, from_be_bytes, from_le_bytes);
208    read_scalar!(i8, from_be_bytes, from_le_bytes);
209    read_scalar!(i16, from_be_bytes, from_le_bytes);
210    read_scalar!(i32, from_be_bytes, from_le_bytes);
211    read_scalar!(i64, from_be_bytes, from_le_bytes);
212    read_scalar!(i128, from_be_bytes, from_le_bytes);
213    read_scalar!(f32, from_be_bytes, from_le_bytes);
214    read_scalar!(f64, from_be_bytes, from_le_bytes);
215}
216
217impl<A: AsyncRead + ?Sized> AsyncReadExt for A {}
218
219/// Implemented as an extension trait, adding utility methods to all
220/// [`AsyncReadAt`] types. Callers will tend to import this trait instead of
221/// [`AsyncReadAt`].
222pub trait AsyncReadAtExt: AsyncReadAt {
223    /// Read the exact number of bytes required to fill `buffer`.
224    ///
225    /// This function reads as many bytes as necessary to completely fill the
226    /// uninitialized space of specified `buffer`.
227    ///
228    /// # Errors
229    ///
230    /// If this function encounters an "end of file" before completely filling
231    /// the buffer, it returns an error of the kind
232    /// [`ErrorKind::UnexpectedEof`]. The contents of `buffer` are unspecified
233    /// in this case.
234    ///
235    /// If any other read error is encountered then this function immediately
236    /// returns. The contents of `buffer` are unspecified in this case.
237    ///
238    /// If this function returns an error, it is unspecified how many bytes it
239    /// has read, but it will never read more than would be necessary to
240    /// completely fill the buffer.
241    ///
242    /// [`ErrorKind::UnexpectedEof`]: std::io::ErrorKind::UnexpectedEof
243    async fn read_exact_at<T: IoBufMut>(&self, mut buf: T, pos: u64) -> BufResult<(), T> {
244        loop_read_exact!(
245            buf,
246            buf.buf_capacity(),
247            read,
248            loop self.read_at(buf.slice(read..), pos + read as u64)
249        );
250    }
251
252    /// Read all bytes as [`String`] until EOF in this source, placing them into
253    /// `buffer`.
254    async fn read_to_string_at(&mut self, buf: String, pos: u64) -> BufResult<usize, String> {
255        let BufResult(res, buf) = self.read_to_end_at(buf.into_bytes(), pos).await;
256        after_read_to_string(res, buf)
257    }
258
259    /// Read all bytes until EOF in this source, placing them into `buffer`.
260    ///
261    /// All bytes read from this source will be appended to the specified buffer
262    /// `buffer`. This function will continuously call [`read_at()`] to append
263    /// more data to `buffer` until [`read_at()`] returns [`Ok(0)`].
264    ///
265    /// If successful, this function will return the total number of bytes read.
266    ///
267    /// [`read_at()`]: AsyncReadAt::read_at
268    async fn read_to_end_at<#[cfg(feature = "allocator_api")] A: Allocator + 'static>(
269        &self,
270        mut buffer: t_alloc!(Vec, u8, A),
271        pos: u64,
272    ) -> BufResult<usize, t_alloc!(Vec, u8, A)> {
273        loop_read_to_end!(buffer, total: u64, loop self.read_at(buffer.slice(total as usize..), pos + total))
274    }
275
276    /// Like [`AsyncReadExt::read_vectored_exact`], expect that it reads at a
277    /// specified position.
278    async fn read_vectored_exact_at<T: IoVectoredBufMut>(
279        &self,
280        mut buf: T,
281        pos: u64,
282    ) -> BufResult<(), T> {
283        let len = buf.total_capacity();
284        loop_read_exact!(buf, len, read, loop self.read_vectored_at(buf.slice_mut(read), pos + read as u64));
285    }
286}
287
288impl<A: AsyncReadAt + ?Sized> AsyncReadAtExt for A {}
compio_io/read/ext.rs

compio_io/read/
ext.rs
