Skip to main content

moirai_async/io/
compat.rs

1// These are used only by the `tokio-compat` trait-bridge impls below.
2#[cfg(feature = "tokio-compat")]
3use std::io::{self, IoSlice};
4#[cfg(feature = "tokio-compat")]
5use std::pin::Pin;
6#[cfg(feature = "tokio-compat")]
7use std::task::{Context, Poll};
8
9#[cfg(feature = "tokio-compat")]
10use crate::io::traits::{AsyncBufRead, AsyncRead, AsyncWrite};
11
12#[cfg(feature = "tokio-compat")]
13use tokio_dep as tokio;
14
15#[cfg(all(test, feature = "tokio-compat"))]
16mod tests;
17
18/// Exposes a Moirai reader, writer, or buffered reader through the
19/// `tokio::io` traits.
20///
21/// The wrapper forwards every poll, including the task context, so the waker a
22/// Tokio caller registers is the waker the Moirai type stores. It keeps the
23/// layout of `T` and allocates nothing.
24#[repr(transparent)]
25pub struct TokioCompat<T> {
26    inner: T,
27}
28
29impl<T> TokioCompat<T> {
30    /// Asserts at compile time that the wrapper adds no size or padding.
31    const TRANSPARENT: () = assert!(
32        size_of::<Self>() == size_of::<T>() && align_of::<Self>() == align_of::<T>(),
33        "compatibility wrapper must keep the layout of its inner type",
34    );
35
36    /// Create a new Tokio compatibility wrapper.
37    pub fn new(inner: T) -> Self {
38        let () = Self::TRANSPARENT;
39        Self { inner }
40    }
41
42    /// Extract the inner type.
43    pub fn into_inner(self) -> T {
44        self.inner
45    }
46}
47
48impl<T> From<T> for TokioCompat<T> {
49    fn from(inner: T) -> Self {
50        Self::new(inner)
51    }
52}
53
54/// Exposes a `tokio::io` reader, writer, or buffered reader through the Moirai
55/// I/O traits.
56///
57/// The wrapper forwards every poll, including the task context, so the waker a
58/// Moirai caller registers is the waker the Tokio type stores. It keeps the
59/// layout of `T` and allocates nothing.
60#[repr(transparent)]
61pub struct MoiraiCompat<T> {
62    inner: T,
63}
64
65impl<T> MoiraiCompat<T> {
66    /// Asserts at compile time that the wrapper adds no size or padding.
67    const TRANSPARENT: () = assert!(
68        size_of::<Self>() == size_of::<T>() && align_of::<Self>() == align_of::<T>(),
69        "compatibility wrapper must keep the layout of its inner type",
70    );
71
72    /// Create a new Moirai compatibility wrapper.
73    pub fn new(inner: T) -> Self {
74        let () = Self::TRANSPARENT;
75        Self { inner }
76    }
77
78    /// Extract the inner type.
79    pub fn into_inner(self) -> T {
80        self.inner
81    }
82}
83
84impl<T> From<T> for MoiraiCompat<T> {
85    fn from(inner: T) -> Self {
86        Self::new(inner)
87    }
88}
89
90#[cfg(feature = "tokio-compat")]
91impl<T: AsyncRead + Unpin> tokio::io::AsyncRead for TokioCompat<T> {
92    fn poll_read(
93        mut self: Pin<&mut Self>,
94        cx: &mut Context<'_>,
95        buf: &mut tokio::io::ReadBuf<'_>,
96    ) -> Poll<io::Result<()>> {
97        let unfilled = buf.initialize_unfilled();
98        match Pin::new(&mut self.inner).poll_read(cx, unfilled) {
99            Poll::Ready(Ok(n)) => {
100                buf.advance(n);
101                Poll::Ready(Ok(()))
102            }
103            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
104            Poll::Pending => Poll::Pending,
105        }
106    }
107}
108
109#[cfg(feature = "tokio-compat")]
110impl<T: AsyncWrite + Unpin> tokio::io::AsyncWrite for TokioCompat<T> {
111    fn poll_write(
112        mut self: Pin<&mut Self>,
113        cx: &mut Context<'_>,
114        buf: &[u8],
115    ) -> Poll<io::Result<usize>> {
116        Pin::new(&mut self.inner).poll_write(cx, buf)
117    }
118
119    fn poll_write_vectored(
120        mut self: Pin<&mut Self>,
121        cx: &mut Context<'_>,
122        bufs: &[IoSlice<'_>],
123    ) -> Poll<io::Result<usize>> {
124        Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
125    }
126
127    fn is_write_vectored(&self) -> bool {
128        self.inner.is_write_vectored()
129    }
130
131    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
132        Pin::new(&mut self.inner).poll_flush(cx)
133    }
134
135    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
136        Pin::new(&mut self.inner).poll_shutdown(cx)
137    }
138}
139
140#[cfg(feature = "tokio-compat")]
141impl<T: AsyncBufRead + Unpin> tokio::io::AsyncBufRead for TokioCompat<T> {
142    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
143        Pin::new(&mut self.get_mut().inner).poll_fill_buf(cx)
144    }
145
146    fn consume(self: Pin<&mut Self>, amt: usize) {
147        Pin::new(&mut self.get_mut().inner).consume(amt);
148    }
149}
150
151#[cfg(feature = "tokio-compat")]
152impl<T: tokio::io::AsyncRead + Unpin> AsyncRead for MoiraiCompat<T> {
153    fn poll_read(
154        mut self: Pin<&mut Self>,
155        cx: &mut Context<'_>,
156        buf: &mut [u8],
157    ) -> Poll<io::Result<usize>> {
158        let mut read_buf = tokio::io::ReadBuf::new(buf);
159        match Pin::new(&mut self.inner).poll_read(cx, &mut read_buf) {
160            Poll::Ready(Ok(())) => Poll::Ready(Ok(read_buf.filled().len())),
161            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
162            Poll::Pending => Poll::Pending,
163        }
164    }
165}
166
167#[cfg(feature = "tokio-compat")]
168impl<T: tokio::io::AsyncWrite + Unpin> AsyncWrite for MoiraiCompat<T> {
169    fn poll_write(
170        mut self: Pin<&mut Self>,
171        cx: &mut Context<'_>,
172        buf: &[u8],
173    ) -> Poll<io::Result<usize>> {
174        Pin::new(&mut self.inner).poll_write(cx, buf)
175    }
176
177    fn poll_write_vectored(
178        mut self: Pin<&mut Self>,
179        cx: &mut Context<'_>,
180        bufs: &[IoSlice<'_>],
181    ) -> Poll<io::Result<usize>> {
182        Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
183    }
184
185    fn is_write_vectored(&self) -> bool {
186        self.inner.is_write_vectored()
187    }
188
189    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
190        Pin::new(&mut self.inner).poll_flush(cx)
191    }
192
193    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
194        Pin::new(&mut self.inner).poll_shutdown(cx)
195    }
196}
197
198#[cfg(feature = "tokio-compat")]
199impl<T: tokio::io::AsyncBufRead + Unpin> AsyncBufRead for MoiraiCompat<T> {
200    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
201        Pin::new(&mut self.get_mut().inner).poll_fill_buf(cx)
202    }
203
204    fn consume(self: Pin<&mut Self>, amt: usize) {
205        Pin::new(&mut self.get_mut().inner).consume(amt);
206    }
207}