1 | use crate::io::AsyncWrite; |
2 | |
3 | use pin_project_lite::pin_project; |
4 | use std::future::Future; |
5 | use std::io; |
6 | use std::marker::PhantomPinned; |
7 | use std::mem; |
8 | use std::pin::Pin; |
9 | use std::task::{Context, Poll}; |
10 | |
11 | pin_project! { |
12 | #[derive(Debug)] |
13 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
14 | pub struct WriteAll<'a, W: ?Sized> { |
15 | writer: &'a mut W, |
16 | buf: &'a [u8], |
17 | // Make this future `!Unpin` for compatibility with async trait methods. |
18 | #[pin] |
19 | _pin: PhantomPinned, |
20 | } |
21 | } |
22 | |
23 | pub(crate) fn write_all<'a, W>(writer: &'a mut W, buf: &'a [u8]) -> WriteAll<'a, W> |
24 | where |
25 | W: AsyncWrite + Unpin + ?Sized, |
26 | { |
27 | WriteAll { |
28 | writer, |
29 | buf, |
30 | _pin: PhantomPinned, |
31 | } |
32 | } |
33 | |
34 | impl<W> Future for WriteAll<'_, W> |
35 | where |
36 | W: AsyncWrite + Unpin + ?Sized, |
37 | { |
38 | type Output = io::Result<()>; |
39 | |
40 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
41 | let me = self.project(); |
42 | while !me.buf.is_empty() { |
43 | let n = ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf))?; |
44 | { |
45 | let (_, rest) = mem::take(&mut *me.buf).split_at(n); |
46 | *me.buf = rest; |
47 | } |
48 | if n == 0 { |
49 | return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); |
50 | } |
51 | } |
52 | |
53 | Poll::Ready(Ok(())) |
54 | } |
55 | } |
56 | |