1 use crate::io::AsyncWrite; 2 3 use bytes::BufMut; 4 use pin_project_lite::pin_project; 5 use std::future::Future; 6 use std::io; 7 use std::mem::size_of; 8 use std::pin::Pin; 9 use std::task::{Context, Poll}; 10 11 macro_rules! writer { 12 ($name:ident, $ty:ty, $writer:ident) => { 13 writer!($name, $ty, $writer, size_of::<$ty>()); 14 }; 15 ($name:ident, $ty:ty, $writer:ident, $bytes:expr) => { 16 pin_project! { 17 #[doc(hidden)] 18 pub struct $name<W> { 19 #[pin] 20 dst: W, 21 buf: [u8; $bytes], 22 written: u8, 23 } 24 } 25 26 impl<W> $name<W> { 27 pub(crate) fn new(w: W, value: $ty) -> Self { 28 let mut writer = $name { 29 buf: [0; $bytes], 30 written: 0, 31 dst: w, 32 }; 33 BufMut::$writer(&mut &mut writer.buf[..], value); 34 writer 35 } 36 } 37 38 impl<W> Future for $name<W> 39 where 40 W: AsyncWrite, 41 { 42 type Output = io::Result<()>; 43 44 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 45 let mut me = self.project(); 46 47 if *me.written == $bytes as u8 { 48 return Poll::Ready(Ok(())); 49 } 50 51 while *me.written < $bytes as u8 { 52 *me.written += match me 53 .dst 54 .as_mut() 55 .poll_write(cx, &me.buf[*me.written as usize..]) 56 { 57 Poll::Pending => return Poll::Pending, 58 Poll::Ready(Err(e)) => return Poll::Ready(Err(e.into())), 59 Poll::Ready(Ok(0)) => { 60 return Poll::Ready(Err(io::ErrorKind::WriteZero.into())); 61 } 62 Poll::Ready(Ok(n)) => n as u8, 63 }; 64 } 65 Poll::Ready(Ok(())) 66 } 67 } 68 }; 69 } 70 71 macro_rules! writer8 { 72 ($name:ident, $ty:ty) => { 73 pin_project! { 74 #[doc(hidden)] 75 pub struct $name<W> { 76 #[pin] 77 dst: W, 78 byte: $ty, 79 } 80 } 81 82 impl<W> $name<W> { 83 pub(crate) fn new(dst: W, byte: $ty) -> Self { 84 Self { dst, byte } 85 } 86 } 87 88 impl<W> Future for $name<W> 89 where 90 W: AsyncWrite, 91 { 92 type Output = io::Result<()>; 93 94 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { 95 let me = self.project(); 96 97 let buf = [*me.byte as u8]; 98 99 match me.dst.poll_write(cx, &buf[..]) { 100 Poll::Pending => Poll::Pending, 101 Poll::Ready(Err(e)) => Poll::Ready(Err(e.into())), 102 Poll::Ready(Ok(0)) => Poll::Ready(Err(io::ErrorKind::WriteZero.into())), 103 Poll::Ready(Ok(1)) => Poll::Ready(Ok(())), 104 Poll::Ready(Ok(_)) => unreachable!(), 105 } 106 } 107 } 108 }; 109 } 110 111 writer8!(WriteU8, u8); 112 writer8!(WriteI8, i8); 113 114 writer!(WriteU16, u16, put_u16); 115 writer!(WriteU32, u32, put_u32); 116 writer!(WriteU64, u64, put_u64); 117 writer!(WriteU128, u128, put_u128); 118 119 writer!(WriteI16, i16, put_i16); 120 writer!(WriteI32, i32, put_i32); 121 writer!(WriteI64, i64, put_i64); 122 writer!(WriteI128, i128, put_i128); 123 124 writer!(WriteU16Le, u16, put_u16_le); 125 writer!(WriteU32Le, u32, put_u32_le); 126 writer!(WriteU64Le, u64, put_u64_le); 127 writer!(WriteU128Le, u128, put_u128_le); 128 129 writer!(WriteI16Le, i16, put_i16_le); 130 writer!(WriteI32Le, i32, put_i32_le); 131 writer!(WriteI64Le, i64, put_i64_le); 132 writer!(WriteI128Le, i128, put_i128_le); 133