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)) => return Poll::Ready(Err(e.into())),
61                        Poll::Ready(Ok(())) => {
62                            let n = buf.filled().len();
63                            if n == 0 {
64                                return Poll::Ready(Err(UnexpectedEof.into()));
65                            }
66
67                            n as u8
68                        }
69                    };
70                }
71
72                let num = Buf::$reader(&mut &me.buf[..]);
73
74                Poll::Ready(Ok(num))
75            }
76        }
77    };
78}
79
80macro_rules! reader8 {
81    ($name:ident, $ty:ty) => {
82        pin_project! {
83            /// Future returned from `read_u8`
84            #[doc(hidden)]
85            #[must_use = "futures do nothing unless you `.await` or poll them"]
86            pub struct $name<R> {
87                #[pin]
88                reader: R,
89                // Make this future `!Unpin` for compatibility with async trait methods.
90                #[pin]
91                _pin: PhantomPinned,
92            }
93        }
94
95        impl<R> $name<R> {
96            pub(crate) fn new(reader: R) -> $name<R> {
97                $name {
98                    reader,
99                    _pin: PhantomPinned,
100                }
101            }
102        }
103
104        impl<R> Future for $name<R>
105        where
106            R: AsyncRead,
107        {
108            type Output = io::Result<$ty>;
109
110            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
111                let me = self.project();
112
113                let mut buf = [0; 1];
114                let mut buf = ReadBuf::new(&mut buf);
115                match me.reader.poll_read(cx, &mut buf) {
116                    Poll::Pending => Poll::Pending,
117                    Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())),
118                    Poll::Ready(Ok(())) => {
119                        if buf.filled().len() == 0 {
120                            return Poll::Ready(Err(UnexpectedEof.into()));
121                        }
122
123                        Poll::Ready(Ok(buf.filled()[0] as $ty))
124                    }
125                }
126            }
127        }
128    };
129}
130
131reader8!(ReadU8, u8);
132reader8!(ReadI8, i8);
133
134reader!(ReadU16, u16, get_u16);
135reader!(ReadU32, u32, get_u32);
136reader!(ReadU64, u64, get_u64);
137reader!(ReadU128, u128, get_u128);
138
139reader!(ReadI16, i16, get_i16);
140reader!(ReadI32, i32, get_i32);
141reader!(ReadI64, i64, get_i64);
142reader!(ReadI128, i128, get_i128);
143
144reader!(ReadF32, f32, get_f32);
145reader!(ReadF64, f64, get_f64);
146
147reader!(ReadU16Le, u16, get_u16_le);
148reader!(ReadU32Le, u32, get_u32_le);
149reader!(ReadU64Le, u64, get_u64_le);
150reader!(ReadU128Le, u128, get_u128_le);
151
152reader!(ReadI16Le, i16, get_i16_le);
153reader!(ReadI32Le, i32, get_i32_le);
154reader!(ReadI64Le, i64, get_i64_le);
155reader!(ReadI128Le, i128, get_i128_le);
156
157reader!(ReadF32Le, f32, get_f32_le);
158reader!(ReadF64Le, f64, get_f64_le);