1 use crate::io::AsyncWrite;
2 
3 use bytes::BufMut;
4 use pin_project_lite::pin_project;
5 use std::future::Future;
6 use std::io;
7 use std::marker::PhantomPinned;
8 use std::mem::size_of;
9 use std::pin::Pin;
10 use std::task::{Context, Poll};
11 
12 macro_rules! writer {
13     ($name:ident, $ty:ty, $writer:ident) => {
14         writer!($name, $ty, $writer, size_of::<$ty>());
15     };
16     ($name:ident, $ty:ty, $writer:ident, $bytes:expr) => {
17         pin_project! {
18             #[doc(hidden)]
19             #[must_use = "futures do nothing unless you `.await` or poll them"]
20             pub struct $name<W> {
21                 #[pin]
22                 dst: W,
23                 buf: [u8; $bytes],
24                 written: u8,
25                 // Make this future `!Unpin` for compatibility with async trait methods.
26                 #[pin]
27                 _pin: PhantomPinned,
28             }
29         }
30 
31         impl<W> $name<W> {
32             pub(crate) fn new(w: W, value: $ty) -> Self {
33                 let mut writer = Self {
34                     buf: [0; $bytes],
35                     written: 0,
36                     dst: w,
37                     _pin: PhantomPinned,
38                 };
39                 BufMut::$writer(&mut &mut writer.buf[..], value);
40                 writer
41             }
42         }
43 
44         impl<W> Future for $name<W>
45         where
46             W: AsyncWrite,
47         {
48             type Output = io::Result<()>;
49 
50             fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
51                 let mut me = self.project();
52 
53                 if *me.written == $bytes as u8 {
54                     return Poll::Ready(Ok(()));
55                 }
56 
57                 while *me.written < $bytes as u8 {
58                     *me.written += match me
59                         .dst
60                         .as_mut()
61                         .poll_write(cx, &me.buf[*me.written as usize..])
62                     {
63                         Poll::Pending => return Poll::Pending,
64                         Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())),
65                         Poll::Ready(Ok(0)) => {
66                             return Poll::Ready(Err(io::ErrorKind::WriteZero.into()));
67                         }
68                         Poll::Ready(Ok(n)) => n as u8,
69                     };
70                 }
71                 Poll::Ready(Ok(()))
72             }
73         }
74     };
75 }
76 
77 macro_rules! writer8 {
78     ($name:ident, $ty:ty) => {
79         pin_project! {
80             #[doc(hidden)]
81             #[must_use = "futures do nothing unless you `.await` or poll them"]
82             pub struct $name<W> {
83                 #[pin]
84                 dst: W,
85                 byte: $ty,
86                 // Make this future `!Unpin` for compatibility with async trait methods.
87                 #[pin]
88                 _pin: PhantomPinned,
89             }
90         }
91 
92         impl<W> $name<W> {
93             pub(crate) fn new(dst: W, byte: $ty) -> Self {
94                 Self {
95                     dst,
96                     byte,
97                     _pin: PhantomPinned,
98                 }
99             }
100         }
101 
102         impl<W> Future for $name<W>
103         where
104             W: AsyncWrite,
105         {
106             type Output = io::Result<()>;
107 
108             fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
109                 let me = self.project();
110 
111                 let buf = [*me.byte as u8];
112 
113                 match me.dst.poll_write(cx, &buf[..]) {
114                     Poll::Pending => Poll::Pending,
115                     Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
116                     Poll::Ready(Ok(0)) => Poll::Ready(Err(io::ErrorKind::WriteZero.into())),
117                     Poll::Ready(Ok(1)) => Poll::Ready(Ok(())),
118                     Poll::Ready(Ok(_)) => unreachable!(),
119                 }
120             }
121         }
122     };
123 }
124 
125 writer8!(WriteU8, u8);
126 writer8!(WriteI8, i8);
127 
128 writer!(WriteU16, u16, put_u16);
129 writer!(WriteU32, u32, put_u32);
130 writer!(WriteU64, u64, put_u64);
131 writer!(WriteU128, u128, put_u128);
132 
133 writer!(WriteI16, i16, put_i16);
134 writer!(WriteI32, i32, put_i32);
135 writer!(WriteI64, i64, put_i64);
136 writer!(WriteI128, i128, put_i128);
137 
138 writer!(WriteU16Le, u16, put_u16_le);
139 writer!(WriteU32Le, u32, put_u32_le);
140 writer!(WriteU64Le, u64, put_u64_le);
141 writer!(WriteU128Le, u128, put_u128_le);
142 
143 writer!(WriteI16Le, i16, put_i16_le);
144 writer!(WriteI32Le, i32, put_i32_le);
145 writer!(WriteI64Le, i64, put_i64_le);
146 writer!(WriteI128Le, i128, put_i128_le);
147