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::mem::size_of;
9use std::pin::Pin;
10use std::task::{Context, Poll};
11
12macro_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
77macro_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
125writer8!(WriteU8, u8);
126writer8!(WriteI8, i8);
127
128writer!(WriteU16, u16, put_u16);
129writer!(WriteU32, u32, put_u32);
130writer!(WriteU64, u64, put_u64);
131writer!(WriteU128, u128, put_u128);
132
133writer!(WriteI16, i16, put_i16);
134writer!(WriteI32, i32, put_i32);
135writer!(WriteI64, i64, put_i64);
136writer!(WriteI128, i128, put_i128);
137
138writer!(WriteF32, f32, put_f32);
139writer!(WriteF64, f64, put_f64);
140
141writer!(WriteU16Le, u16, put_u16_le);
142writer!(WriteU32Le, u32, put_u32_le);
143writer!(WriteU64Le, u64, put_u64_le);
144writer!(WriteU128Le, u128, put_u128_le);
145
146writer!(WriteI16Le, i16, put_i16_le);
147writer!(WriteI32Le, i32, put_i32_le);
148writer!(WriteI64Le, i64, put_i64_le);
149writer!(WriteI128Le, i128, put_i128_le);
150
151writer!(WriteF32Le, f32, put_f32_le);
152writer!(WriteF64Le, f64, put_f64_le);
153