compio_io/write/
ext.rs

1use compio_buf::{BufResult, IntoInner, IoBuf, IoVectoredBuf};
2
3use crate::{AsyncWrite, AsyncWriteAt, IoResult};
4
5/// Shared code for write a scalar value into the underlying writer.
6macro_rules! write_scalar {
7    ($t:ty, $be:ident, $le:ident) => {
8        ::paste::paste! {
9            #[doc = concat!("Write a big endian `", stringify!($t), "` into the underlying writer.")]
10            async fn [< write_ $t >](&mut self, num: $t) -> IoResult<()> {
11                use ::compio_buf::{arrayvec::ArrayVec, BufResult};
12
13                const LEN: usize = ::std::mem::size_of::<$t>();
14                let BufResult(res, _) = self
15                    .write_all(ArrayVec::<u8, LEN>::from(num.$be()))
16                    .await;
17                res
18            }
19
20            #[doc = concat!("Write a little endian `", stringify!($t), "` into the underlying writer.")]
21            async fn [< write_ $t _le >](&mut self, num: $t) -> IoResult<()> {
22                use ::compio_buf::{arrayvec::ArrayVec, BufResult};
23
24                const LEN: usize = ::std::mem::size_of::<$t>();
25                let BufResult(res, _) = self
26                    .write_all(ArrayVec::<u8, LEN>::from(num.$le()))
27                    .await;
28                res
29            }
30        }
31    };
32}
33
34/// Shared code for loop writing until all contents are written.
35macro_rules! loop_write_all {
36    ($buf:ident, $len:expr, $needle:ident,loop $expr_expr:expr) => {
37        let len = $len;
38        let mut $needle = 0;
39
40        while $needle < len {
41            match $expr_expr.await.into_inner() {
42                BufResult(Ok(0), buf) => {
43                    return BufResult(
44                        Err(::std::io::Error::new(
45                            ::std::io::ErrorKind::WriteZero,
46                            "failed to write whole buffer",
47                        )),
48                        buf,
49                    );
50                }
51                BufResult(Ok(n), buf) => {
52                    $needle += n;
53                    $buf = buf;
54                }
55                BufResult(Err(ref e), buf) if e.kind() == ::std::io::ErrorKind::Interrupted => {
56                    $buf = buf;
57                }
58                BufResult(Err(e), buf) => return BufResult(Err(e), buf),
59            }
60        }
61
62        return BufResult(Ok(()), $buf);
63    };
64}
65
66macro_rules! loop_write_vectored {
67    ($buf:ident, $iter:ident, $read_expr:expr) => {{
68        let mut $iter = match $buf.owned_iter() {
69            Ok(buf) => buf,
70            Err(buf) => return BufResult(Ok(0), buf),
71        };
72
73        loop {
74            if $iter.buf_len() > 0 {
75                return $read_expr.await.into_inner();
76            }
77
78            match $iter.next() {
79                Ok(next) => $iter = next,
80                Err(buf) => return BufResult(Ok(0), buf),
81            }
82        }
83    }};
84}
85
86/// Implemented as an extension trait, adding utility methods to all
87/// [`AsyncWrite`] types. Callers will tend to import this trait instead of
88/// [`AsyncWrite`].
89pub trait AsyncWriteExt: AsyncWrite {
90    /// Creates a "by reference" adaptor for this instance of [`AsyncWrite`].
91    ///
92    /// The returned adapter also implements [`AsyncWrite`] and will simply
93    /// borrow this current writer.
94    fn by_ref(&mut self) -> &mut Self
95    where
96        Self: Sized,
97    {
98        self
99    }
100
101    /// Write the entire contents of a buffer into this writer.
102    async fn write_all<T: IoBuf>(&mut self, mut buf: T) -> BufResult<(), T> {
103        loop_write_all!(
104            buf,
105            buf.buf_len(),
106            needle,
107            loop self.write(buf.slice(needle..))
108        );
109    }
110
111    /// Write the entire contents of a buffer into this writer. Like
112    /// [`AsyncWrite::write_vectored`], except that it tries to write the entire
113    /// contents of the buffer into this writer.
114    async fn write_vectored_all<T: IoVectoredBuf>(&mut self, mut buf: T) -> BufResult<(), T> {
115        let len = buf.total_len();
116        loop_write_all!(buf, len, needle, loop self.write_vectored(buf.slice(needle)));
117    }
118
119    write_scalar!(u8, to_be_bytes, to_le_bytes);
120    write_scalar!(u16, to_be_bytes, to_le_bytes);
121    write_scalar!(u32, to_be_bytes, to_le_bytes);
122    write_scalar!(u64, to_be_bytes, to_le_bytes);
123    write_scalar!(u128, to_be_bytes, to_le_bytes);
124    write_scalar!(i8, to_be_bytes, to_le_bytes);
125    write_scalar!(i16, to_be_bytes, to_le_bytes);
126    write_scalar!(i32, to_be_bytes, to_le_bytes);
127    write_scalar!(i64, to_be_bytes, to_le_bytes);
128    write_scalar!(i128, to_be_bytes, to_le_bytes);
129    write_scalar!(f32, to_be_bytes, to_le_bytes);
130    write_scalar!(f64, to_be_bytes, to_le_bytes);
131}
132
133impl<A: AsyncWrite + ?Sized> AsyncWriteExt for A {}
134
135/// Implemented as an extension trait, adding utility methods to all
136/// [`AsyncWriteAt`] types. Callers will tend to import this trait instead of
137/// [`AsyncWriteAt`].
138pub trait AsyncWriteAtExt: AsyncWriteAt {
139    /// Like [`AsyncWriteAt::write_at`], except that it tries to write the
140    /// entire contents of the buffer into this writer.
141    async fn write_all_at<T: IoBuf>(&mut self, mut buf: T, pos: u64) -> BufResult<(), T> {
142        loop_write_all!(
143            buf,
144            buf.buf_len(),
145            needle,
146            loop self.write_at(buf.slice(needle..), pos + needle as u64)
147        );
148    }
149
150    /// Like [`AsyncWriteAt::write_vectored_at`], expect that it tries to write
151    /// the entire contents of the buffer into this writer.
152    async fn write_vectored_all_at<T: IoVectoredBuf>(
153        &mut self,
154        mut buf: T,
155        pos: u64,
156    ) -> BufResult<(), T> {
157        let len = buf.total_len();
158        loop_write_all!(buf, len, needle, loop self.write_vectored_at(buf.slice(needle), pos + needle as u64));
159    }
160}
161
162impl<A: AsyncWriteAt + ?Sized> AsyncWriteAtExt for A {}
compio_io/write/ext.rs

compio_io/write/
ext.rs
