use core::pin::Pin;
use core::task::{Context, Poll};
#[cfg(feature = "alloc")]
use alloc::vec::Vec;
pub trait AsyncRead {
type Error;
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, Self::Error>>;
}
pub trait AsyncWrite {
type Error;
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Self::Error>>;
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
}
#[cfg(feature = "alloc")]
pub struct Prepend<S> {
leftover: Vec<u8>,
cursor: usize,
inner: S,
}
#[cfg(feature = "alloc")]
impl<S> Prepend<S> {
#[must_use]
pub fn new(leftover: Vec<u8>, inner: S) -> Self {
Self {
leftover,
cursor: 0,
inner,
}
}
pub fn into_inner(self) -> S {
self.inner
}
}
#[cfg(feature = "alloc")]
impl<S: AsyncRead + Unpin> AsyncRead for Prepend<S> {
type Error = S::Error;
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, Self::Error>> {
let this = self.get_mut();
if this.cursor < this.leftover.len() {
let remaining = &this.leftover[this.cursor..];
let count = remaining.len().min(buf.len());
buf[..count].copy_from_slice(&remaining[..count]);
this.cursor += count;
if this.cursor == this.leftover.len() {
this.leftover = Vec::new();
}
return Poll::Ready(Ok(count));
}
Pin::new(&mut this.inner).poll_read(cx, buf)
}
}
#[cfg(feature = "alloc")]
impl<S: AsyncWrite + Unpin> AsyncWrite for Prepend<S> {
type Error = S::Error;
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Self::Error>> {
Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.get_mut().inner).poll_close(cx)
}
}
#[cfg(feature = "io-async-compat")]
pub struct FromFutures<S>(pub S);
#[cfg(feature = "io-async-compat")]
impl<S: futures_io::AsyncRead + Unpin> AsyncRead for FromFutures<S> {
type Error = std::io::Error;
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, Self::Error>> {
Pin::new(&mut self.get_mut().0).poll_read(cx, buf)
}
}
#[cfg(feature = "io-async-compat")]
impl<S: futures_io::AsyncWrite + Unpin> AsyncWrite for FromFutures<S> {
type Error = std::io::Error;
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Self::Error>> {
Pin::new(&mut self.get_mut().0).poll_write(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.get_mut().0).poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.get_mut().0).poll_close(cx)
}
}
#[cfg(feature = "io-async-compat")]
pub struct IntoFutures<S>(pub S);
#[cfg(feature = "io-async-compat")]
impl<S> futures_io::AsyncRead for IntoFutures<S>
where
S: AsyncRead + Unpin,
S::Error: Into<std::io::Error>,
{
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, std::io::Error>> {
match Pin::new(&mut self.get_mut().0).poll_read(cx, buf) {
Poll::Ready(Ok(count)) => Poll::Ready(Ok(count)),
Poll::Ready(Err(error)) => Poll::Ready(Err(error.into())),
Poll::Pending => Poll::Pending,
}
}
}
#[cfg(feature = "io-async-compat")]
impl<S> futures_io::AsyncWrite for IntoFutures<S>
where
S: AsyncWrite + Unpin,
S::Error: Into<std::io::Error>,
{
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, std::io::Error>> {
match Pin::new(&mut self.get_mut().0).poll_write(cx, buf) {
Poll::Ready(Ok(count)) => Poll::Ready(Ok(count)),
Poll::Ready(Err(error)) => Poll::Ready(Err(error.into())),
Poll::Pending => Poll::Pending,
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
match Pin::new(&mut self.get_mut().0).poll_flush(cx) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
Poll::Ready(Err(error)) => Poll::Ready(Err(error.into())),
Poll::Pending => Poll::Pending,
}
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), std::io::Error>> {
match Pin::new(&mut self.get_mut().0).poll_close(cx) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
Poll::Ready(Err(error)) => Poll::Ready(Err(error.into())),
Poll::Pending => Poll::Pending,
}
}
}
#[cfg(feature = "io-async-compat-tokio")]
pub struct FromTokio<S>(pub S);
#[cfg(feature = "io-async-compat-tokio")]
impl<S: tokio::io::AsyncRead + Unpin> AsyncRead for FromTokio<S> {
type Error = std::io::Error;
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, Self::Error>> {
let mut read_buf = tokio::io::ReadBuf::new(buf);
match Pin::new(&mut self.get_mut().0).poll_read(cx, &mut read_buf) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(read_buf.filled().len())),
Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
Poll::Pending => Poll::Pending,
}
}
}
#[cfg(feature = "io-async-compat-tokio")]
impl<S: tokio::io::AsyncWrite + Unpin> AsyncWrite for FromTokio<S> {
type Error = std::io::Error;
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Self::Error>> {
Pin::new(&mut self.get_mut().0).poll_write(cx, buf)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.get_mut().0).poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Pin::new(&mut self.get_mut().0).poll_shutdown(cx)
}
}
#[cfg(feature = "io-async-compat-tokio")]
pub struct IntoTokio<S>(pub S);
#[cfg(feature = "io-async-compat-tokio")]
impl<S> tokio::io::AsyncRead for IntoTokio<S>
where
S: AsyncRead + Unpin,
S::Error: Into<std::io::Error>,
{
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
let dst = buf.initialize_unfilled();
match Pin::new(&mut self.get_mut().0).poll_read(cx, dst) {
Poll::Ready(Ok(count)) => {
buf.advance(count);
Poll::Ready(Ok(()))
}
Poll::Ready(Err(error)) => Poll::Ready(Err(error.into())),
Poll::Pending => Poll::Pending,
}
}
}
#[cfg(feature = "io-async-compat-tokio")]
impl<S> tokio::io::AsyncWrite for IntoTokio<S>
where
S: AsyncWrite + Unpin,
S::Error: Into<std::io::Error>,
{
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<std::io::Result<usize>> {
match Pin::new(&mut self.get_mut().0).poll_write(cx, buf) {
Poll::Ready(Ok(count)) => Poll::Ready(Ok(count)),
Poll::Ready(Err(error)) => Poll::Ready(Err(error.into())),
Poll::Pending => Poll::Pending,
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
match Pin::new(&mut self.get_mut().0).poll_flush(cx) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
Poll::Ready(Err(error)) => Poll::Ready(Err(error.into())),
Poll::Pending => Poll::Pending,
}
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
match Pin::new(&mut self.get_mut().0).poll_close(cx) {
Poll::Ready(Ok(())) => Poll::Ready(Ok(())),
Poll::Ready(Err(error)) => Poll::Ready(Err(error.into())),
Poll::Pending => Poll::Pending,
}
}
}
#[cfg(test)]
mod tests {
use core::task::Waker;
#[cfg(feature = "io-async-compat-tokio")]
use tokio::io::{AsyncRead as _, AsyncWrite as _};
use super::*;
#[derive(Debug, Clone, Copy, PartialEq)]
enum MockError {
Broken,
}
struct SliceReader<'a> {
data: &'a [u8],
pos: usize,
}
impl<'a> SliceReader<'a> {
fn new(data: &'a [u8]) -> Self {
Self { data, pos: 0 }
}
}
impl AsyncRead for SliceReader<'_> {
type Error = MockError;
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, Self::Error>> {
let this = self.get_mut();
let remaining = &this.data[this.pos..];
let count = remaining.len().min(buf.len());
buf[..count].copy_from_slice(&remaining[..count]);
this.pos += count;
Poll::Ready(Ok(count))
}
}
struct BrokenReader;
impl AsyncRead for BrokenReader {
type Error = MockError;
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
_buf: &mut [u8],
) -> Poll<Result<usize, Self::Error>> {
Poll::Ready(Err(MockError::Broken))
}
}
struct PendingOnceReader<'a> {
inner: SliceReader<'a>,
pending_remaining: u8,
}
impl AsyncRead for PendingOnceReader<'_> {
type Error = MockError;
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<Result<usize, Self::Error>> {
let this = self.get_mut();
if this.pending_remaining > 0 {
this.pending_remaining -= 1;
cx.waker().wake_by_ref();
return Poll::Pending;
}
Pin::new(&mut this.inner).poll_read(cx, buf)
}
}
struct SliceWriter<'a> {
buf: &'a mut [u8],
pos: usize,
}
impl AsyncWrite for SliceWriter<'_> {
type Error = MockError;
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
data: &[u8],
) -> Poll<Result<usize, Self::Error>> {
let this = self.get_mut();
let remaining = this.buf.len() - this.pos;
let count = remaining.min(data.len());
this.buf[this.pos..this.pos + count].copy_from_slice(&data[..count]);
this.pos += count;
Poll::Ready(Ok(count))
}
fn poll_flush(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
#[test]
fn read_happy_path_fills_buffer_in_one_poll() {
let data = b"proxima";
let mut reader = SliceReader::new(data);
let mut context = Context::from_waker(Waker::noop());
let mut buf = [0u8; 7];
let outcome = Pin::new(&mut reader).poll_read(&mut context, &mut buf);
assert_eq!(
outcome,
Poll::Ready(Ok(7)),
"a buffer sized to fit must drain the source in one poll"
);
assert_eq!(&buf, data, "bytes read must match source data");
}
#[test]
fn read_returns_zero_at_eof() {
let data = b"ab";
let mut reader = SliceReader::new(data);
let mut context = Context::from_waker(Waker::noop());
let mut buf = [0u8; 2];
let _ = Pin::new(&mut reader).poll_read(&mut context, &mut buf);
let eof_outcome = Pin::new(&mut reader).poll_read(&mut context, &mut buf);
assert_eq!(
eof_outcome,
Poll::Ready(Ok(0)),
"a drained reader must signal eof, not error or pend"
);
}
#[test]
fn read_reassembles_across_undersized_buffers() {
let data = b"proxima-tier-lift";
let mut reader = SliceReader::new(data);
let mut context = Context::from_waker(Waker::noop());
let mut assembled = [0u8; 32];
let mut written = 0usize;
loop {
let mut chunk = [0u8; 3];
match Pin::new(&mut reader).poll_read(&mut context, &mut chunk) {
Poll::Ready(Ok(0)) => break,
Poll::Ready(Ok(count)) => {
assembled[written..written + count].copy_from_slice(&chunk[..count]);
written += count;
}
Poll::Ready(Err(error)) => panic!("unexpected read error: {error:?}"),
Poll::Pending => panic!("SliceReader never pends"),
}
}
assert_eq!(
written,
data.len(),
"partial reads must sum to the full source length"
);
assert_eq!(
&assembled[..written],
data,
"reassembled bytes must match source data byte-for-byte"
);
}
#[test]
fn read_error_propagates_through_poll() {
let mut reader = BrokenReader;
let mut context = Context::from_waker(Waker::noop());
let mut buf = [0u8; 4];
let outcome = Pin::new(&mut reader).poll_read(&mut context, &mut buf);
assert_eq!(
outcome,
Poll::Ready(Err(MockError::Broken)),
"broken reader must surface its error, not silently succeed"
);
}
#[cfg(feature = "alloc")]
#[test]
fn prepend_replays_leftover_then_delegates_to_inner() {
let mut prepended = Prepend::new(b"PRI ".to_vec(), SliceReader::new(b"proxima"));
let mut context = Context::from_waker(Waker::noop());
let mut assembled: Vec<u8> = Vec::new();
loop {
let mut chunk = [0u8; 4];
match Pin::new(&mut prepended).poll_read(&mut context, &mut chunk) {
Poll::Ready(Ok(0)) => break,
Poll::Ready(Ok(count)) => assembled.extend_from_slice(&chunk[..count]),
other => panic!("prepend read must not error or pend here: {other:?}"),
}
}
assert_eq!(
&assembled[..],
b"PRI proxima",
"the prepended leftover must precede the inner reader's bytes, intact"
);
}
#[cfg(feature = "io-async-compat")]
#[test]
fn from_futures_bridge_feeds_the_prepend_adapter() {
let socket = futures::io::Cursor::new(b"proxima".to_vec());
let mut adapter = Prepend::new(b">>".to_vec(), FromFutures(socket));
let mut context = Context::from_waker(Waker::noop());
let mut assembled: Vec<u8> = Vec::new();
loop {
let mut chunk = [0u8; 4];
match Pin::new(&mut adapter).poll_read(&mut context, &mut chunk) {
Poll::Ready(Ok(0)) => break,
Poll::Ready(Ok(count)) => assembled.extend_from_slice(&chunk[..count]),
other => panic!("bridged read must not error or pend here: {other:?}"),
}
}
assert_eq!(
&assembled[..],
b">>proxima",
"prepend leftover then the bridged futures reader's bytes, in order"
);
}
#[cfg(feature = "io-async-compat-tokio")]
struct TokioSliceReader<'a> {
data: &'a [u8],
pos: usize,
}
#[cfg(feature = "io-async-compat-tokio")]
impl<'a> TokioSliceReader<'a> {
fn new(data: &'a [u8]) -> Self {
Self { data, pos: 0 }
}
}
#[cfg(feature = "io-async-compat-tokio")]
impl tokio::io::AsyncRead for TokioSliceReader<'_> {
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &mut tokio::io::ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
let this = self.get_mut();
let remaining = &this.data[this.pos..];
let count = remaining.len().min(buf.remaining());
buf.put_slice(&remaining[..count]);
this.pos += count;
Poll::Ready(Ok(()))
}
}
#[cfg(feature = "io-async-compat-tokio")]
#[test]
fn from_tokio_bridge_feeds_the_prepend_adapter() {
let socket = TokioSliceReader::new(b"proxima");
let mut adapter = Prepend::new(b">>".to_vec(), FromTokio(socket));
let mut context = Context::from_waker(Waker::noop());
let mut assembled: Vec<u8> = Vec::new();
loop {
let mut chunk = [0u8; 4];
match Pin::new(&mut adapter).poll_read(&mut context, &mut chunk) {
Poll::Ready(Ok(0)) => break,
Poll::Ready(Ok(count)) => assembled.extend_from_slice(&chunk[..count]),
other => panic!("bridged read must not error or pend here: {other:?}"),
}
}
assert_eq!(
&assembled[..],
b">>proxima",
"prepend leftover then the bridged tokio reader's bytes, in order"
);
}
#[cfg(feature = "io-async-compat-tokio")]
#[test]
fn into_tokio_round_trips_through_from_tokio() {
let data = b"proxima-tier-lift";
let socket = TokioSliceReader::new(data);
let mut bridged = IntoTokio(FromTokio(socket));
let mut context = Context::from_waker(Waker::noop());
let mut backing = [0u8; 32];
let mut read_buf = tokio::io::ReadBuf::new(&mut backing);
let outcome = Pin::new(&mut bridged).poll_read(&mut context, &mut read_buf);
assert!(
matches!(outcome, Poll::Ready(Ok(()))),
"round-trip read must not error or pend: {outcome:?}"
);
assert_eq!(
read_buf.filled(),
data,
"IntoTokio(FromTokio(_)) must preserve bytes byte-for-byte"
);
}
#[cfg(feature = "io-async-compat-tokio")]
#[test]
fn into_tokio_write_flush_shutdown_delegate_to_from_tokio() {
let sink: Vec<u8> = Vec::new();
let mut bridged = IntoTokio(FromTokio(sink));
let mut context = Context::from_waker(Waker::noop());
let write_outcome = Pin::new(&mut bridged).poll_write(&mut context, b"proxima!");
assert!(
matches!(write_outcome, Poll::Ready(Ok(8))),
"writer must accept the full payload: {write_outcome:?}"
);
let flush_outcome = Pin::new(&mut bridged).poll_flush(&mut context);
assert!(
matches!(flush_outcome, Poll::Ready(Ok(()))),
"flush must complete synchronously for a Vec sink: {flush_outcome:?}"
);
let shutdown_outcome = Pin::new(&mut bridged).poll_shutdown(&mut context);
assert!(
matches!(shutdown_outcome, Poll::Ready(Ok(()))),
"shutdown must complete synchronously for a Vec sink: {shutdown_outcome:?}"
);
assert_eq!(
&bridged.0.0[..],
b"proxima!",
"bytes actually written must match the source payload"
);
}
#[test]
fn write_flush_close_happy_path() {
let mut backing = [0u8; 8];
let mut context = Context::from_waker(Waker::noop());
{
let mut writer = SliceWriter {
buf: &mut backing,
pos: 0,
};
let write_outcome = Pin::new(&mut writer).poll_write(&mut context, b"proxima!");
assert_eq!(
write_outcome,
Poll::Ready(Ok(8)),
"writer must accept the full payload"
);
let flush_outcome = Pin::new(&mut writer).poll_flush(&mut context);
assert_eq!(
flush_outcome,
Poll::Ready(Ok(())),
"flush must complete synchronously for an in-memory sink"
);
let close_outcome = Pin::new(&mut writer).poll_close(&mut context);
assert_eq!(
close_outcome,
Poll::Ready(Ok(())),
"close must complete synchronously for an in-memory sink"
);
}
assert_eq!(
&backing, b"proxima!",
"bytes actually written must match the source payload"
);
}
#[test]
fn pending_registers_waker_then_makes_progress() {
let data = b"waking";
let mut reader = PendingOnceReader {
inner: SliceReader::new(data),
pending_remaining: 1,
};
let mut context = Context::from_waker(Waker::noop());
let mut buf = [0u8; 6];
let first_poll = Pin::new(&mut reader).poll_read(&mut context, &mut buf);
assert_eq!(
first_poll,
Poll::Pending,
"reader not yet ready must return pending and register the waker"
);
let second_poll = Pin::new(&mut reader).poll_read(&mut context, &mut buf);
assert_eq!(
second_poll,
Poll::Ready(Ok(6)),
"reader must make progress on the next poll once ready"
);
assert_eq!(
&buf, data,
"bytes read after pending must match source data"
);
}
}