#[cfg(feature = "allocator_api")]
use std::alloc::Allocator;
use std::io::Cursor;
use compio_buf::{BufResult, IntoInner, IoBuf, IoVectoredBuf, buf_try, t_alloc};
use crate::IoResult;
mod buf;
#[macro_use]
mod ext;
pub use buf::*;
pub use ext::*;
pub trait AsyncWrite {
async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T>;
async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
loop_write_vectored!(buf, total: usize, n, iter, loop self.write(iter),
break if n == 0 || n < iter.buf_len() {
Some(Ok(total))
} else {
None
}
)
}
async fn flush(&mut self) -> IoResult<()>;
async fn shutdown(&mut self) -> IoResult<()>;
}
impl<A: AsyncWrite + ?Sized> AsyncWrite for &mut A {
async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
(**self).write(buf).await
}
async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
(**self).write_vectored(buf).await
}
async fn flush(&mut self) -> IoResult<()> {
(**self).flush().await
}
async fn shutdown(&mut self) -> IoResult<()> {
(**self).shutdown().await
}
}
impl<W: AsyncWrite + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator> AsyncWrite
for t_alloc!(Box, W, A)
{
async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
(**self).write(buf).await
}
async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
(**self).write_vectored(buf).await
}
async fn flush(&mut self) -> IoResult<()> {
(**self).flush().await
}
async fn shutdown(&mut self) -> IoResult<()> {
(**self).shutdown().await
}
}
impl AsyncWrite for Vec<u8> {
async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
self.extend_from_slice(buf.as_slice());
BufResult(Ok(buf.buf_len()), buf)
}
async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
let len = buf.iter_buf().map(|b| b.buf_len()).sum();
self.reserve(len - self.len());
for buf in buf.iter_buf() {
self.extend_from_slice(buf.as_slice());
}
BufResult(Ok(len), buf)
}
async fn flush(&mut self) -> IoResult<()> {
Ok(())
}
async fn shutdown(&mut self) -> IoResult<()> {
Ok(())
}
}
pub trait AsyncWriteAt {
async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T>;
async fn write_vectored_at<T: IoVectoredBuf>(
&mut self,
buf: T,
pos: u64,
) -> BufResult<usize, T> {
loop_write_vectored!(buf, total: u64, n, iter, loop self.write_at(iter, pos + total),
break if n == 0 || n < iter.buf_len() {
Some(Ok(total as usize))
} else {
None
}
)
}
}
impl<A: AsyncWriteAt + ?Sized> AsyncWriteAt for &mut A {
async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
(**self).write_at(buf, pos).await
}
async fn write_vectored_at<T: IoVectoredBuf>(
&mut self,
buf: T,
pos: u64,
) -> BufResult<usize, T> {
(**self).write_vectored_at(buf, pos).await
}
}
impl<W: AsyncWriteAt + ?Sized, #[cfg(feature = "allocator_api")] A: Allocator> AsyncWriteAt
for t_alloc!(Box, W, A)
{
async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
(**self).write_at(buf, pos).await
}
async fn write_vectored_at<T: IoVectoredBuf>(
&mut self,
buf: T,
pos: u64,
) -> BufResult<usize, T> {
(**self).write_vectored_at(buf, pos).await
}
}
macro_rules! impl_write_at {
($($(const $len:ident =>)? $ty:ty),*) => {
$(
impl<$(const $len: usize)?> AsyncWriteAt for $ty {
async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
let pos = (pos as usize).min(self.len());
let slice = buf.as_slice();
let n = slice.len().min(self.len() - pos);
self[pos..pos + n].copy_from_slice(&slice[..n]);
BufResult(Ok(n), buf)
}
}
)*
}
}
impl_write_at!([u8], const LEN => [u8; LEN]);
impl<#[cfg(feature = "allocator_api")] A: Allocator> AsyncWriteAt for t_alloc!(Vec, u8, A) {
async fn write_at<T: IoBuf>(&mut self, buf: T, pos: u64) -> BufResult<usize, T> {
let pos = pos as usize;
let slice = buf.as_slice();
if pos <= self.len() {
let n = slice.len().min(self.len() - pos);
if n < slice.len() {
self.reserve(slice.len() - n);
self[pos..pos + n].copy_from_slice(&slice[..n]);
self.extend_from_slice(&slice[n..]);
} else {
self[pos..pos + n].copy_from_slice(slice);
}
BufResult(Ok(n), buf)
} else {
self.reserve(pos - self.len() + slice.len());
self.resize(pos, 0);
self.extend_from_slice(slice);
BufResult(Ok(slice.len()), buf)
}
}
}
impl<A: AsyncWriteAt> AsyncWrite for Cursor<A> {
async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
let pos = self.position();
let (n, buf) = buf_try!(self.get_mut().write_at(buf, pos).await);
self.set_position(pos + n as u64);
BufResult(Ok(n), buf)
}
async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
let pos = self.position();
let (n, buf) = buf_try!(self.get_mut().write_vectored_at(buf, pos).await);
self.set_position(pos + n as u64);
BufResult(Ok(n), buf)
}
async fn flush(&mut self) -> IoResult<()> {
Ok(())
}
async fn shutdown(&mut self) -> IoResult<()> {
Ok(())
}
}