compio_io\read/
ext.rs

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

compio_io\read/
ext.rs
