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