1use std::io::{self, IoSlice};
4use std::ops::DerefMut;
5use std::pin::Pin;
6use std::task::{Context, Poll};
7
8pub trait AsyncWrite {
10 fn poll_write(
12 self: Pin<&mut Self>,
13 cx: &mut Context<'_>,
14 buf: &[u8],
15 ) -> Poll<io::Result<usize>>;
16
17 fn poll_write_vectored(
19 self: Pin<&mut Self>,
20 cx: &mut Context<'_>,
21 bufs: &[IoSlice<'_>],
22 ) -> Poll<io::Result<usize>> {
23 for buf in bufs {
25 if !buf.is_empty() {
26 return self.poll_write(cx, buf);
27 }
28 }
29 Poll::Ready(Ok(0))
30 }
31
32 fn is_write_vectored(&self) -> bool {
34 false
35 }
36
37 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>>;
39
40 fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>>;
42}
43
44pub trait AsyncWriteVectored: AsyncWrite {
46 fn poll_write_vectored(
48 self: Pin<&mut Self>,
49 cx: &mut Context<'_>,
50 bufs: &[IoSlice<'_>],
51 ) -> Poll<io::Result<usize>> {
52 AsyncWrite::poll_write_vectored(self, cx, bufs)
53 }
54
55 fn is_write_vectored(&self) -> bool {
57 AsyncWrite::is_write_vectored(self)
58 }
59}
60
61impl<W> AsyncWriteVectored for W where W: AsyncWrite + ?Sized {}
62
63impl AsyncWrite for Vec<u8> {
64 fn poll_write(
65 self: Pin<&mut Self>,
66 _cx: &mut Context<'_>,
67 buf: &[u8],
68 ) -> Poll<io::Result<usize>> {
69 let this = self.get_mut();
70 this.extend_from_slice(buf);
71 Poll::Ready(Ok(buf.len()))
72 }
73
74 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
75 Poll::Ready(Ok(()))
76 }
77
78 fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
79 Poll::Ready(Ok(()))
80 }
81}
82
83impl AsyncWrite for std::io::Cursor<&mut [u8]> {
84 fn poll_write(
85 self: Pin<&mut Self>,
86 _cx: &mut Context<'_>,
87 buf: &[u8],
88 ) -> Poll<io::Result<usize>> {
89 use std::io::Write as _;
90
91 let this = self.get_mut();
92 let n = this.write(buf)?;
93 Poll::Ready(Ok(n))
94 }
95
96 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
97 Poll::Ready(Ok(()))
98 }
99
100 fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
101 Poll::Ready(Ok(()))
102 }
103}
104
105impl AsyncWrite for std::io::Cursor<Vec<u8>> {
106 fn poll_write(
107 self: Pin<&mut Self>,
108 _cx: &mut Context<'_>,
109 buf: &[u8],
110 ) -> Poll<io::Result<usize>> {
111 use std::io::Write as _;
112
113 let this = self.get_mut();
114 let n = this.write(buf)?;
115 Poll::Ready(Ok(n))
116 }
117
118 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
119 Poll::Ready(Ok(()))
120 }
121
122 fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
123 Poll::Ready(Ok(()))
124 }
125}
126
127impl AsyncWrite for std::io::Cursor<Box<[u8]>> {
128 fn poll_write(
129 self: Pin<&mut Self>,
130 _cx: &mut Context<'_>,
131 buf: &[u8],
132 ) -> Poll<io::Result<usize>> {
133 use std::io::Write as _;
134
135 let this = self.get_mut();
136 let n = this.write(buf)?;
137 Poll::Ready(Ok(n))
138 }
139
140 fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
141 Poll::Ready(Ok(()))
142 }
143
144 fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
145 Poll::Ready(Ok(()))
146 }
147}
148
149impl<W> AsyncWrite for &mut W
150where
151 W: AsyncWrite + Unpin + ?Sized,
152{
153 fn poll_write(
154 self: Pin<&mut Self>,
155 cx: &mut Context<'_>,
156 buf: &[u8],
157 ) -> Poll<io::Result<usize>> {
158 let this = self.get_mut();
159 Pin::new(&mut **this).poll_write(cx, buf)
160 }
161
162 fn poll_write_vectored(
163 self: Pin<&mut Self>,
164 cx: &mut Context<'_>,
165 bufs: &[IoSlice<'_>],
166 ) -> Poll<io::Result<usize>> {
167 let this = self.get_mut();
168 Pin::new(&mut **this).poll_write_vectored(cx, bufs)
169 }
170
171 fn is_write_vectored(&self) -> bool {
172 (**self).is_write_vectored()
173 }
174
175 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
176 let this = self.get_mut();
177 Pin::new(&mut **this).poll_flush(cx)
178 }
179
180 fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
181 let this = self.get_mut();
182 Pin::new(&mut **this).poll_shutdown(cx)
183 }
184}
185
186impl<W> AsyncWrite for Box<W>
187where
188 W: AsyncWrite + Unpin + ?Sized,
189{
190 fn poll_write(
191 self: Pin<&mut Self>,
192 cx: &mut Context<'_>,
193 buf: &[u8],
194 ) -> Poll<io::Result<usize>> {
195 let this = self.get_mut();
196 Pin::new(&mut **this).poll_write(cx, buf)
197 }
198
199 fn poll_write_vectored(
200 self: Pin<&mut Self>,
201 cx: &mut Context<'_>,
202 bufs: &[IoSlice<'_>],
203 ) -> Poll<io::Result<usize>> {
204 let this = self.get_mut();
205 Pin::new(&mut **this).poll_write_vectored(cx, bufs)
206 }
207
208 fn is_write_vectored(&self) -> bool {
209 (**self).is_write_vectored()
210 }
211
212 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
213 let this = self.get_mut();
214 Pin::new(&mut **this).poll_flush(cx)
215 }
216
217 fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
218 let this = self.get_mut();
219 Pin::new(&mut **this).poll_shutdown(cx)
220 }
221}
222
223impl<W, P> AsyncWrite for Pin<P>
224where
225 P: DerefMut<Target = W> + Unpin,
226 W: AsyncWrite + Unpin + ?Sized,
227{
228 fn poll_write(
229 self: Pin<&mut Self>,
230 cx: &mut Context<'_>,
231 buf: &[u8],
232 ) -> Poll<io::Result<usize>> {
233 let this = self.get_mut();
234 Pin::new(&mut **this).poll_write(cx, buf)
235 }
236
237 fn poll_write_vectored(
238 self: Pin<&mut Self>,
239 cx: &mut Context<'_>,
240 bufs: &[IoSlice<'_>],
241 ) -> Poll<io::Result<usize>> {
242 let this = self.get_mut();
243 Pin::new(&mut **this).poll_write_vectored(cx, bufs)
244 }
245
246 fn is_write_vectored(&self) -> bool {
247 (**self).is_write_vectored()
248 }
249
250 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
251 let this = self.get_mut();
252 Pin::new(&mut **this).poll_flush(cx)
253 }
254
255 fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
256 let this = self.get_mut();
257 Pin::new(&mut **this).poll_shutdown(cx)
258 }
259}
260
261#[cfg(test)]
262mod tests {
263 use super::*;
264 use std::sync::Arc;
265 use std::task::{Context, Wake, Waker};
266
267 struct NoopWaker;
268
269 impl Wake for NoopWaker {
270 fn wake(self: Arc<Self>) {}
271 }
272
273 fn noop_waker() -> Waker {
274 Waker::from(Arc::new(NoopWaker))
275 }
276
277 fn init_test(name: &str) {
278 crate::test_utils::init_test_logging();
279 crate::test_phase!(name);
280 }
281
282 #[test]
283 fn write_to_vec() {
284 init_test("write_to_vec");
285 let mut output = Vec::new();
286 let waker = noop_waker();
287 let mut cx = Context::from_waker(&waker);
288
289 let poll = Pin::new(&mut output).poll_write(&mut cx, b"hello");
290 let ready = matches!(poll, Poll::Ready(Ok(5)));
291 crate::assert_with_log!(ready, "write 5", true, ready);
292 crate::assert_with_log!(output == b"hello", "output", b"hello", output);
293 crate::test_complete!("write_to_vec");
294 }
295
296 #[test]
297 fn write_to_cursor() {
298 init_test("write_to_cursor");
299 let mut buf = [0u8; 8];
300 let mut cursor = std::io::Cursor::new(&mut buf[..]);
301 let waker = noop_waker();
302 let mut cx = Context::from_waker(&waker);
303
304 let poll = Pin::new(&mut cursor).poll_write(&mut cx, b"test");
305 let ready = matches!(poll, Poll::Ready(Ok(4)));
306 crate::assert_with_log!(ready, "write 4", true, ready);
307 crate::assert_with_log!(&buf[..4] == b"test", "buf", b"test", &buf[..4]);
308 crate::test_complete!("write_to_cursor");
309 }
310
311 #[test]
312 fn flush_and_shutdown_vec() {
313 init_test("flush_and_shutdown_vec");
314 let mut output = Vec::new();
315 let waker = noop_waker();
316 let mut cx = Context::from_waker(&waker);
317
318 let poll = Pin::new(&mut output).poll_flush(&mut cx);
319 let ready = matches!(poll, Poll::Ready(Ok(())));
320 crate::assert_with_log!(ready, "flush ready", true, ready);
321
322 let poll = Pin::new(&mut output).poll_shutdown(&mut cx);
323 let ready = matches!(poll, Poll::Ready(Ok(())));
324 crate::assert_with_log!(ready, "shutdown ready", true, ready);
325 crate::test_complete!("flush_and_shutdown_vec");
326 }
327
328 #[test]
329 fn write_via_ref() {
330 init_test("write_via_ref");
331 let mut output = Vec::new();
332 let mut writer = &mut output;
333 let waker = noop_waker();
334 let mut cx = Context::from_waker(&waker);
335
336 let poll = Pin::new(&mut writer).poll_write(&mut cx, b"via ref");
337 let ready = matches!(poll, Poll::Ready(Ok(7)));
338 crate::assert_with_log!(ready, "write 7", true, ready);
339 crate::assert_with_log!(output == b"via ref", "output", b"via ref", output);
340 crate::test_complete!("write_via_ref");
341 }
342
343 #[test]
344 fn write_via_box() {
345 init_test("write_via_box");
346 let mut output: Box<Vec<u8>> = Box::default();
347 let waker = noop_waker();
348 let mut cx = Context::from_waker(&waker);
349
350 let poll = Pin::new(&mut output).poll_write(&mut cx, b"boxed");
351 let ready = matches!(poll, Poll::Ready(Ok(5)));
352 crate::assert_with_log!(ready, "write 5", true, ready);
353 crate::assert_with_log!(*output == b"boxed", "output", b"boxed", *output);
354 crate::test_complete!("write_via_box");
355 }
356}