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