tokio/io/util/
read_int.rs

1use crate::io::{AsyncRead, ReadBuf};
2
3use bytes::Buf;
4use pin_project_lite::pin_project;
5use std::future::Future;
6use std::io;
7use std::io::ErrorKind::UnexpectedEof;
8use std::marker::PhantomPinned;
9use std::pin::Pin;
10use std::task::{Context, Poll};
11
12macro_rules! reader {
13    ($name:ident, $ty:ty, $reader:ident) => {
14        reader!($name, $ty, $reader, std::mem::size_of::<$ty>());
15    };
16    ($name:ident, $ty:ty, $reader: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<R> {
21                #[pin]
22                src: R,
23                buf: [u8; $bytes],
24                read: u8,
25                // Make this future `!Unpin` for compatibility with async trait methods.
26                #[pin]
27                _pin: PhantomPinned,
28            }
29        }
30
31        impl<R> $name<R> {
32            pub(crate) fn new(src: R) -> Self {
33                $name {
34                    src,
35                    buf: [0; $bytes],
36                    read: 0,
37                    _pin: PhantomPinned,
38                }
39            }
40        }
41
42        impl<R> Future for $name<R>
43        where
44            R: AsyncRead,
45        {
46            type Output = io::Result<$ty>;
47
48            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
49                let mut me = self.project();
50
51                if *me.read == $bytes as u8 {
52                    return Poll::Ready(Ok(Buf::$reader(&mut &me.buf[..])));
53                }
54
55                while *me.read < $bytes as u8 {
56                    let mut buf = ReadBuf::new(&mut me.buf[*me.read as usize..]);
57
58                    *me.read += match me.src.as_mut().poll_read(cx, &mut buf) {
59                        Poll::Pending => return Poll::Pending,
60                        Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue,
61                        Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())),
62                        Poll::Ready(Ok(())) => {
63                            let n = buf.filled().len();
64                            if n == 0 {
65                                return Poll::Ready(Err(UnexpectedEof.into()));
66                            }
67
68                            n as u8
69                        }
70                    };
71                }
72
73                let num = Buf::$reader(&mut &me.buf[..]);
74
75                Poll::Ready(Ok(num))
76            }
77        }
78    };
79}
80
81macro_rules! reader8 {
82    ($name:ident, $ty:ty) => {
83        pin_project! {
84            /// Future returned from `read_u8`
85            #[doc(hidden)]
86            #[must_use = "futures do nothing unless you `.await` or poll them"]
87            pub struct $name<R> {
88                #[pin]
89                reader: R,
90                // Make this future `!Unpin` for compatibility with async trait methods.
91                #[pin]
92                _pin: PhantomPinned,
93            }
94        }
95
96        impl<R> $name<R> {
97            pub(crate) fn new(reader: R) -> $name<R> {
98                $name {
99                    reader,
100                    _pin: PhantomPinned,
101                }
102            }
103        }
104
105        impl<R> Future for $name<R>
106        where
107            R: AsyncRead,
108        {
109            type Output = io::Result<$ty>;
110
111            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
112                let mut me = self.project();
113
114                let mut buf = [0; 1];
115                let mut buf = ReadBuf::new(&mut buf);
116                loop {
117                    match me.reader.as_mut().poll_read(cx, &mut buf) {
118                        Poll::Pending => return Poll::Pending,
119                        Poll::Ready(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue,
120                        Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())),
121                        Poll::Ready(Ok(())) => {
122                            if buf.filled().len() == 0 {
123                                return Poll::Ready(Err(UnexpectedEof.into()));
124                            }
125
126                            return Poll::Ready(Ok(buf.filled()[0] as $ty));
127                        }
128                    }
129                }
130            }
131        }
132    };
133}
134
135reader8!(ReadU8, u8);
136reader8!(ReadI8, i8);
137
138reader!(ReadU16, u16, get_u16);
139reader!(ReadU32, u32, get_u32);
140reader!(ReadU64, u64, get_u64);
141reader!(ReadU128, u128, get_u128);
142
143reader!(ReadI16, i16, get_i16);
144reader!(ReadI32, i32, get_i32);
145reader!(ReadI64, i64, get_i64);
146reader!(ReadI128, i128, get_i128);
147
148reader!(ReadF32, f32, get_f32);
149reader!(ReadF64, f64, get_f64);
150
151reader!(ReadU16Le, u16, get_u16_le);
152reader!(ReadU32Le, u32, get_u32_le);
153reader!(ReadU64Le, u64, get_u64_le);
154reader!(ReadU128Le, u128, get_u128_le);
155
156reader!(ReadI16Le, i16, get_i16_le);
157reader!(ReadI32Le, i32, get_i32_le);
158reader!(ReadI64Le, i64, get_i64_le);
159reader!(ReadI128Le, i128, get_i128_le);
160
161reader!(ReadF32Le, f32, get_f32_le);
162reader!(ReadF64Le, f64, get_f64_le);
tokio/io/util/read_int.rs

tokio/io/util/
read_int.rs
