nuclei/io/util/
write_ext.rs

1//! [`AsyncWriteExt`]: the writes every [`AsyncWrite`] gets, as futures.
2
3use core::future::Future;
4use core::pin::Pin;
5use core::task::{Context, Poll, ready};
6use std::io::{self, IoSlice};
7
8use futures_io::AsyncWrite;
9
10/// One typed-write method per line: name, type, conversion, what it writes.
11macro_rules! write_ints {
12    ($($name:ident, $ty:ty, $conv:ident, $what:literal;)+) => {$(
13        #[doc = concat!("Writes ", $what, ", all of its bytes.")]
14        fn $name(&mut self, n: $ty) -> WriteNumber<'_, Self, { core::mem::size_of::<$ty>() }>
15        where
16            Self: Unpin,
17        {
18            WriteNumber {
19                writer: self,
20                bytes: n.$conv(),
21                done: 0,
22            }
23        }
24    )+};
25}
26
27/// Writes on any [`AsyncWrite`]: a slice, all of a slice, typed numbers in
28/// big or little endian; and its flush and shutdown.
29pub trait AsyncWriteExt: AsyncWrite {
30    /// Writes some of `buf`, answering how many bytes.
31    fn write<'a>(&'a mut self, buf: &'a [u8]) -> Write<'a, Self>
32    where
33        Self: Unpin,
34    {
35        Write { writer: self, buf }
36    }
37
38    /// Writes some of the slices `bufs`, in order, answering how many bytes.
39    fn write_vectored<'a, 'b>(&'a mut self, bufs: &'a [IoSlice<'b>]) -> WriteVectored<'a, 'b, Self>
40    where
41        Self: Unpin,
42    {
43        WriteVectored { writer: self, bufs }
44    }
45
46    /// Writes all of `buf`; a writer that takes no more bytes is a
47    /// [`WriteZero`](io::ErrorKind::WriteZero) error.
48    fn write_all<'a>(&'a mut self, buf: &'a [u8]) -> WriteAll<'a, Self>
49    where
50        Self: Unpin,
51    {
52        WriteAll { writer: self, buf }
53    }
54
55    /// Writes out what the writer holds back, down to its destination.
56    fn flush(&mut self) -> Flush<'_, Self>
57    where
58        Self: Unpin,
59    {
60        Flush { writer: self }
61    }
62
63    /// Flushes and closes the writing side (`poll_close`): a socket sends
64    /// its end of stream, a pipe's reader sees its end.
65    fn shutdown(&mut self) -> Shutdown<'_, Self>
66    where
67        Self: Unpin,
68    {
69        Shutdown { writer: self }
70    }
71
72    write_ints! {
73        write_u8, u8, to_be_bytes, "an unsigned byte";
74        write_i8, i8, to_be_bytes, "a signed byte";
75        write_u16, u16, to_be_bytes, "a big-endian `u16`";
76        write_u16_le, u16, to_le_bytes, "a little-endian `u16`";
77        write_i16, i16, to_be_bytes, "a big-endian `i16`";
78        write_i16_le, i16, to_le_bytes, "a little-endian `i16`";
79        write_u32, u32, to_be_bytes, "a big-endian `u32`";
80        write_u32_le, u32, to_le_bytes, "a little-endian `u32`";
81        write_i32, i32, to_be_bytes, "a big-endian `i32`";
82        write_i32_le, i32, to_le_bytes, "a little-endian `i32`";
83        write_u64, u64, to_be_bytes, "a big-endian `u64`";
84        write_u64_le, u64, to_le_bytes, "a little-endian `u64`";
85        write_i64, i64, to_be_bytes, "a big-endian `i64`";
86        write_i64_le, i64, to_le_bytes, "a little-endian `i64`";
87        write_u128, u128, to_be_bytes, "a big-endian `u128`";
88        write_u128_le, u128, to_le_bytes, "a little-endian `u128`";
89        write_i128, i128, to_be_bytes, "a big-endian `i128`";
90        write_i128_le, i128, to_le_bytes, "a little-endian `i128`";
91        write_f32, f32, to_be_bytes, "a big-endian `f32`";
92        write_f32_le, f32, to_le_bytes, "a little-endian `f32`";
93        write_f64, f64, to_be_bytes, "a big-endian `f64`";
94        write_f64_le, f64, to_le_bytes, "a little-endian `f64`";
95    }
96}
97
98impl<W: AsyncWrite + ?Sized> AsyncWriteExt for W {}
99
100/// The future of [`write`](AsyncWriteExt::write).
101#[derive(Debug)]
102#[must_use = "futures do nothing unless polled"]
103pub struct Write<'a, W: ?Sized> {
104    writer: &'a mut W,
105    buf: &'a [u8],
106}
107
108impl<W: AsyncWrite + Unpin + ?Sized> Future for Write<'_, W> {
109    type Output = io::Result<usize>;
110
111    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
112        let this = self.get_mut();
113        Pin::new(&mut *this.writer).poll_write(cx, this.buf)
114    }
115}
116
117/// The future of [`write_vectored`](AsyncWriteExt::write_vectored).
118#[derive(Debug)]
119#[must_use = "futures do nothing unless polled"]
120pub struct WriteVectored<'a, 'b, W: ?Sized> {
121    writer: &'a mut W,
122    bufs: &'a [IoSlice<'b>],
123}
124
125impl<W: AsyncWrite + Unpin + ?Sized> Future for WriteVectored<'_, '_, W> {
126    type Output = io::Result<usize>;
127
128    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
129        let this = self.get_mut();
130        Pin::new(&mut *this.writer).poll_write_vectored(cx, this.bufs)
131    }
132}
133
134/// The future of [`write_all`](AsyncWriteExt::write_all).
135#[derive(Debug)]
136#[must_use = "futures do nothing unless polled"]
137pub struct WriteAll<'a, W: ?Sized> {
138    writer: &'a mut W,
139    buf: &'a [u8],
140}
141
142impl<W: AsyncWrite + Unpin + ?Sized> Future for WriteAll<'_, W> {
143    type Output = io::Result<()>;
144
145    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
146        let this = self.get_mut();
147        while !this.buf.is_empty() {
148            let n = ready!(Pin::new(&mut *this.writer).poll_write(cx, this.buf))?;
149            if n == 0 {
150                return Poll::Ready(Err(io::ErrorKind::WriteZero.into()));
151            }
152            this.buf = &this.buf[n..];
153        }
154        Poll::Ready(Ok(()))
155    }
156}
157
158/// The future of a typed write ([`write_u32`](AsyncWriteExt::write_u32) and
159/// the like).
160#[derive(Debug)]
161#[must_use = "futures do nothing unless polled"]
162pub struct WriteNumber<'a, W: ?Sized, const N: usize> {
163    writer: &'a mut W,
164    bytes: [u8; N],
165    done: usize,
166}
167
168impl<W: AsyncWrite + Unpin + ?Sized, const N: usize> Future for WriteNumber<'_, W, N> {
169    type Output = io::Result<()>;
170
171    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
172        let this = self.get_mut();
173        while this.done < N {
174            let n = ready!(Pin::new(&mut *this.writer).poll_write(cx, &this.bytes[this.done..]))?;
175            if n == 0 {
176                return Poll::Ready(Err(io::ErrorKind::WriteZero.into()));
177            }
178            this.done += n;
179        }
180        Poll::Ready(Ok(()))
181    }
182}
183
184/// The future of [`flush`](AsyncWriteExt::flush).
185#[derive(Debug)]
186#[must_use = "futures do nothing unless polled"]
187pub struct Flush<'a, W: ?Sized> {
188    writer: &'a mut W,
189}
190
191impl<W: AsyncWrite + Unpin + ?Sized> Future for Flush<'_, W> {
192    type Output = io::Result<()>;
193
194    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
195        Pin::new(&mut *self.get_mut().writer).poll_flush(cx)
196    }
197}
198
199/// The future of [`shutdown`](AsyncWriteExt::shutdown).
200#[derive(Debug)]
201#[must_use = "futures do nothing unless polled"]
202pub struct Shutdown<'a, W: ?Sized> {
203    writer: &'a mut W,
204}
205
206impl<W: AsyncWrite + Unpin + ?Sized> Future for Shutdown<'_, W> {
207    type Output = io::Result<()>;
208
209    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
210        Pin::new(&mut *self.get_mut().writer).poll_close(cx)
211    }
212}
nuclei/io/util/write_ext.rs

nuclei/io/util/
write_ext.rs
