use std::{cmp, io};
use bytes::{Buf, Bytes};
use hyper2::rt::{Read, ReadBufCursor, Write};
use std::{
pin::Pin,
task::{self, Poll},
};
#[derive(Debug)]
pub(crate) struct Rewind<T> {
pub(crate) pre: Option<Bytes>,
pub(crate) inner: T,
}
impl<T> Read for Rewind<T>
where
T: Read + Unpin,
{
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
mut buf: ReadBufCursor<'_>,
) -> Poll<io::Result<()>> {
if let Some(mut prefix) = self.pre.take() {
if !prefix.is_empty() {
let copy_len = cmp::min(prefix.len(), remaining(&mut buf));
put_slice(&mut buf, &prefix[..copy_len]);
prefix.advance(copy_len);
if !prefix.is_empty() {
self.pre = Some(prefix);
}
return Poll::Ready(Ok(()));
}
}
Pin::new(&mut self.inner).poll_read(cx, buf)
}
}
fn remaining(cursor: &mut ReadBufCursor<'_>) -> usize {
unsafe { cursor.as_mut().len() }
}
fn put_slice(cursor: &mut ReadBufCursor<'_>, slice: &[u8]) {
assert!(
remaining(cursor) >= slice.len(),
"buf.len() must fit in remaining()"
);
let amt = slice.len();
unsafe {
cursor.as_mut()[..amt]
.as_mut_ptr()
.cast::<u8>()
.copy_from_nonoverlapping(slice.as_ptr(), amt);
cursor.advance(amt);
}
}
impl<T> Write for Rewind<T>
where
T: Write + Unpin,
{
#[inline(always)]
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.inner).poll_write(cx, buf)
}
fn poll_write_vectored(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
bufs: &[io::IoSlice<'_>],
) -> Poll<io::Result<usize>> {
Pin::new(&mut self.inner).poll_write_vectored(cx, bufs)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.inner).poll_flush(cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.inner).poll_shutdown(cx)
}
fn is_write_vectored(&self) -> bool {
self.inner.is_write_vectored()
}
}