tokio/io/util/
write_int.rs

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)) if e.kind() == io::ErrorKind::Interrupted => continue,
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 mut me = self.project();
110
111                let buf = [*me.byte as u8];
112
113                loop {
114                    match me.dst.as_mut().poll_write(cx, &buf[..]) {
115                        Poll::Pending => return Poll::Pending,
116                        Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue,
117                        Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())),
118                        Poll::Ready(Ok(0)) => {
119                            return Poll::Ready(Err(io::ErrorKind::WriteZero.into()))
120                        }
121                        Poll::Ready(Ok(1)) => return Poll::Ready(Ok(())),
122                        Poll::Ready(Ok(_)) => unreachable!(),
123                    }
124                }
125            }
126        }
127    };
128}
129
130writer8!(WriteU8, u8);
131writer8!(WriteI8, i8);
132
133writer!(WriteU16, u16, put_u16);
134writer!(WriteU32, u32, put_u32);
135writer!(WriteU64, u64, put_u64);
136writer!(WriteU128, u128, put_u128);
137
138writer!(WriteI16, i16, put_i16);
139writer!(WriteI32, i32, put_i32);
140writer!(WriteI64, i64, put_i64);
141writer!(WriteI128, i128, put_i128);
142
143writer!(WriteF32, f32, put_f32);
144writer!(WriteF64, f64, put_f64);
145
146writer!(WriteU16Le, u16, put_u16_le);
147writer!(WriteU32Le, u32, put_u32_le);
148writer!(WriteU64Le, u64, put_u64_le);
149writer!(WriteU128Le, u128, put_u128_le);
150
151writer!(WriteI16Le, i16, put_i16_le);
152writer!(WriteI32Le, i32, put_i32_le);
153writer!(WriteI64Le, i64, put_i64_le);
154writer!(WriteI128Le, i128, put_i128_le);
155
156writer!(WriteF32Le, f32, put_f32_le);
157writer!(WriteF64Le, f64, put_f64_le);
tokio/io/util/write_int.rs

tokio/io/util/
write_int.rs
