use super::DEFAULT_BUF_SIZE;
use futures_core::ready;
use futures_core::task::{Context, Poll};
use futures_io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, IoSlice, SeekFrom};
use pin_project_lite::pin_project;
use std::fmt;
use std::io::{self, Write};
use std::pin::Pin;
use std::ptr;
pin_project! {
pub struct BufWriter<W> {
#[pin]
inner: W,
buf: Vec<u8>,
written: usize,
}
}
impl<W: AsyncWrite> BufWriter<W> {
pub fn new(inner: W) -> Self {
Self::with_capacity(DEFAULT_BUF_SIZE, inner)
}
pub fn with_capacity(cap: usize, inner: W) -> Self {
Self { inner, buf: Vec::with_capacity(cap), written: 0 }
}
pub(super) fn flush_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
let mut this = self.project();
let len = this.buf.len();
let mut ret = Ok(());
while *this.written < len {
match ready!(this.inner.as_mut().poll_write(cx, &this.buf[*this.written..])) {
Ok(0) => {
ret = Err(io::Error::new(
io::ErrorKind::WriteZero,
"failed to write the buffered data",
));
break;
}
Ok(n) => *this.written += n,
Err(e) => {
ret = Err(e);
break;
}
}
}
if *this.written > 0 {
this.buf.drain(..*this.written);
}
*this.written = 0;
Poll::Ready(ret)
}
delegate_access_inner!(inner, W, ());
pub fn buffer(&self) -> &[u8] {
&self.buf
}
pub(super) fn capacity(&self) -> usize {
self.buf.capacity()
}
#[inline]
pub(super) fn spare_capacity(&self) -> usize {
self.buf.capacity() - self.buf.len()
}
pub(super) fn write_to_buf(self: Pin<&mut Self>, buf: &[u8]) -> usize {
let available = self.spare_capacity();
let amt_to_buffer = available.min(buf.len());
unsafe {
self.write_to_buffer_unchecked(&buf[..amt_to_buffer]);
}
amt_to_buffer
}
#[inline]
unsafe fn write_to_buffer_unchecked(self: Pin<&mut Self>, buf: &[u8]) {
debug_assert!(buf.len() <= self.spare_capacity());
let this = self.project();
let old_len = this.buf.len();
let buf_len = buf.len();
let src = buf.as_ptr();
let dst = this.buf.as_mut_ptr().add(old_len);
ptr::copy_nonoverlapping(src, dst, buf_len);
this.buf.set_len(old_len + buf_len);
}
pub(super) fn inner_poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.project().inner.poll_write(cx, buf)
}
pub(super) fn inner_poll_write_vectored(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[IoSlice<'_>],
) -> Poll<io::Result<usize>> {
self.project().inner.poll_write_vectored(cx, bufs)
}
}
impl<W: AsyncWrite> AsyncWrite for BufWriter<W> {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
if self.buf.len() + buf.len() > self.buf.capacity() {
ready!(self.as_mut().flush_buf(cx))?;
}
if buf.len() >= self.buf.capacity() {
self.project().inner.poll_write(cx, buf)
} else {
Poll::Ready(self.project().buf.write(buf))
}
}
fn poll_write_vectored(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
bufs: &[IoSlice<'_>],
) -> Poll<io::Result<usize>> {
let total_len = bufs.iter().map(|b| b.len()).sum::<usize>();
if self.buf.len() + total_len > self.buf.capacity() {
ready!(self.as_mut().flush_buf(cx))?;
}
if total_len >= self.buf.capacity() {
self.project().inner.poll_write_vectored(cx, bufs)
} else {
Poll::Ready(self.project().buf.write_vectored(bufs))
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
ready!(self.as_mut().flush_buf(cx))?;
self.project().inner.poll_flush(cx)
}
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
ready!(self.as_mut().flush_buf(cx))?;
self.project().inner.poll_close(cx)
}
}
impl<W: AsyncRead> AsyncRead for BufWriter<W> {
delegate_async_read!(inner);
}
impl<W: AsyncBufRead> AsyncBufRead for BufWriter<W> {
delegate_async_buf_read!(inner);
}
impl<W: fmt::Debug> fmt::Debug for BufWriter<W> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("BufWriter")
.field("writer", &self.inner)
.field("buffer", &format_args!("{}/{}", self.buf.len(), self.buf.capacity()))
.field("written", &self.written)
.finish()
}
}
impl<W: AsyncWrite + AsyncSeek> AsyncSeek for BufWriter<W> {
fn poll_seek(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
pos: SeekFrom,
) -> Poll<io::Result<u64>> {
ready!(self.as_mut().flush_buf(cx))?;
self.project().inner.poll_seek(cx, pos)
}
}