refactor(io): aware of allocator_api

compio-rs · Mar 3, 2024 · 60a659f · 60a659f
1 parent 1c1b058
commit 60a659f
Show file tree

Hide file tree

Showing 2 changed files with 109 additions and 55 deletions.
diff --git a/compio-io/src/read/mod.rs b/compio-io/src/read/mod.rs
@@ -2,7 +2,9 @@
 use std::alloc::Allocator;
 use std::{io::Cursor, rc::Rc, sync::Arc};
 
-use compio_buf::{buf_try, vec_alloc, BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBufMut};
+use compio_buf::{
+    box_alloc, buf_try, vec_alloc, BufResult, IntoInner, IoBuf, IoBufMut, IoVectoredBufMut,
+};
 
 mod buf;
 #[macro_use]
@@ -59,25 +61,31 @@ pub trait AsyncRead {
     }
 }
 
-macro_rules! impl_read {
-    ($($ty:ty),*) => {
-        $(
-            impl<A: AsyncRead + ?Sized> AsyncRead for $ty {
-                #[inline(always)]
-                async fn read<T: IoBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
-                    (**self).read(buf).await
-                }
+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
-                }
-            }
-        )*
-    };
+    #[inline(always)]
+    async fn read_vectored<T: IoVectoredBufMut>(&mut self, buf: T) -> BufResult<usize, T> {
+        (**self).read_vectored(buf).await
+    }
 }
 
-impl_read!(&mut A, Box<A>);
+impl<R: AsyncRead + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator> AsyncRead
+    for box_alloc!(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]
@@ -139,6 +147,23 @@ macro_rules! impl_read_at {
         )*
     };
 
+    (@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 {
@@ -152,7 +177,8 @@ macro_rules! impl_read_at {
     }
 }
 
-impl_read_at!(@ptr &A, &mut A, Box<A>, Rc<A>, Arc<A>);
+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 vec_alloc!(u8, A) {

diff --git a/compio-io/src/write/mod.rs b/compio-io/src/write/mod.rs
@@ -2,7 +2,7 @@
 use std::alloc::Allocator;
 use std::io::Cursor;
 
-use compio_buf::{buf_try, vec_alloc, BufResult, IntoInner, IoBuf, IoVectoredBuf};
+use compio_buf::{box_alloc, buf_try, vec_alloc, BufResult, IntoInner, IoBuf, IoVectoredBuf};
 
 use crate::IoResult;
 
@@ -50,31 +50,43 @@ pub trait AsyncWrite {
     async fn shutdown(&mut self) -> IoResult<()>;
 }
 
-macro_rules! impl_write {
-    (@ptr $($ty:ty),*) => {
-        $(
-            impl<A: AsyncWrite + ?Sized> AsyncWrite for $ty {
-                async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
-                    (**self).write(buf).await
-                }
+impl<A: AsyncWrite + ?Sized> AsyncWrite for &mut A {
+    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
+        (**self).write(buf).await
+    }
 
-                async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
-                    (**self).write_vectored(buf).await
-                }
+    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
+        (**self).write_vectored(buf).await
+    }
 
-                async fn flush(&mut self) -> IoResult<()> {
-                    (**self).flush().await
-                }
+    async fn flush(&mut self) -> IoResult<()> {
+        (**self).flush().await
+    }
 
-                async fn shutdown(&mut self) -> IoResult<()> {
-                    (**self).shutdown().await
-                }
-            }
-        )*
-    };
+    async fn shutdown(&mut self) -> IoResult<()> {
+        (**self).shutdown().await
+    }
 }
 
-impl_write!(@ptr &mut A, Box<A>);
+impl<W: AsyncWrite + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator> AsyncWrite
+    for box_alloc!(W, A)
+{
+    async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
+        (**self).write(buf).await
+    }
+
+    async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
+        (**self).write_vectored(buf).await
+    }
+
+    async fn flush(&mut self) -> IoResult<()> {
+        (**self).flush().await
+    }
+
+    async fn shutdown(&mut self) -> IoResult<()> {
+        (**self).shutdown().await
+    }
+}
 
 /// Write is implemented for `Vec<u8>` by appending to the vector. The vector
 /// will grow as needed.
@@ -127,21 +139,38 @@ pub trait AsyncWriteAt {
     }
 }
 
-macro_rules! impl_write_at {
-    (@ptr $($ty:ty),*) => {
-        $(
-            impl<A: AsyncWriteAt + ?Sized> AsyncWriteAt for $ty {
-                async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
-                    (**self).write_at(buf, pos).await
-                }
+impl<A: AsyncWriteAt + ?Sized> AsyncWriteAt for &mut A {
+    async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
+        (**self).write_at(buf, pos).await
+    }
 
-                async fn write_vectored_at<T: IoVectoredBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
-                    (**self).write_vectored_at(buf, pos).await
-                }
-            }
-        )*
-    };
-    (@slice $($(const $len:ident =>)? $ty:ty),*) => {
+    async fn write_vectored_at<T: IoVectoredBuf>(
+        &mut self,
+        buf: T,
+        pos: u64,
+    ) -> BufResult<usize, T> {
+        (**self).write_vectored_at(buf, pos).await
+    }
+}
+
+impl<W: AsyncWriteAt + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator> AsyncWriteAt
+    for box_alloc!(W, A)
+{
+    async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
+        (**self).write_at(buf, pos).await
+    }
+
+    async fn write_vectored_at<T: IoVectoredBuf>(
+        &mut self,
+        buf: T,
+        pos: u64,
+    ) -> BufResult<usize, T> {
+        (**self).write_vectored_at(buf, pos).await
+    }
+}
+
+macro_rules! impl_write_at {
+    ($($(const $len:ident =>)? $ty:ty),*) => {
         $(
             impl<$(const $len: usize)?> AsyncWriteAt for $ty {
                 async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
@@ -156,8 +185,7 @@ macro_rules! impl_write_at {
     }
 }
 
-impl_write_at!(@ptr &mut A, Box<A>);
-impl_write_at!(@slice [u8], const LEN => [u8; LEN]);
+impl_write_at!([u8], const LEN => [u8; LEN]);
 
 /// This implementation aligns the behavior of files. If `pos` is larger than
 /// the vector length, the vectored will be extended, and the extended area will