tokio/io/util/
write_int.rs

Help
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);
tokio/io/util/write_int.rs

tokio/io/util/
write_int.rs
