use crate::storage::Storage;
use crate::{Region, RegionMut, Ring};
use alloc::sync::Arc;
use core::fmt;
use core::future::Future;
use core::pin::Pin;
use core::sync::atomic::{AtomicBool, Ordering};
use core::task::{Context, Poll};
use futures::task::AtomicWaker;
pub fn queue_from<T, S>(storage: S) -> (Reader<S>, Writer<S>)
where
S: Storage<T>,
{
let ring = Ring::new(storage.capacity());
queue_from_parts(ring, storage)
}
pub fn queue_from_parts<S>(ring: Ring, storage: S) -> (Reader<S>, Writer<S>) {
let state = Arc::new(State {
ring,
storage,
is_reader_open: AtomicBool::new(true),
is_writer_open: AtomicBool::new(true),
data_available_waker: AtomicWaker::new(),
space_available_waker: AtomicWaker::new(),
});
let reader = Reader {
state: state.clone(),
};
let writer = Writer { state };
(reader, writer)
}
#[cfg(feature = "heap-buffer")]
mod heap_constructors {
use crate::asyncio::{queue_from_parts, Reader, Writer};
use crate::storage::HeapBuffer;
use crate::Ring;
#[cfg_attr(docsrs, doc(cfg(feature = "heap-buffer")))]
pub fn queue<T>(capacity: usize) -> (Reader<HeapBuffer<T>>, Writer<HeapBuffer<T>>)
where
T: Default,
{
let ring = Ring::new(capacity);
let buffer = HeapBuffer::new(capacity);
queue_from_parts(ring, buffer)
}
}
#[cfg(feature = "heap-buffer")]
pub use self::heap_constructors::*;
#[derive(Debug)]
struct State<S> {
ring: Ring,
storage: S,
is_reader_open: AtomicBool,
is_writer_open: AtomicBool,
data_available_waker: AtomicWaker,
space_available_waker: AtomicWaker,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WriteError {
ReaderClosed,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReadExactError {
WriterClosed,
}
#[derive(Debug)]
pub struct Reader<S> {
state: Arc<State<S>>,
}
#[derive(Debug)]
pub struct Writer<S> {
state: Arc<State<S>>,
}
impl<S> State<S> {
fn close_reader(&self) {
let was_open = self.is_reader_open.swap(false, Ordering::AcqRel);
if was_open {
self.space_available_waker.wake();
}
}
fn close_writer(&self) {
let was_open = self.is_writer_open.swap(false, Ordering::AcqRel);
if was_open {
self.data_available_waker.wake();
}
}
}
impl<S> Reader<S> {
#[inline]
pub fn is_writer_open(&self) -> bool {
self.state.is_writer_open.load(Ordering::Acquire)
}
#[inline]
pub fn has_data(&self) -> bool {
let (r0, r1) = self.state.ring.left_ranges();
!r0.is_empty() || !r1.is_empty()
}
#[inline]
pub fn is_full(&self) -> bool {
let (r0, r1) = self.state.ring.right_ranges();
r0.is_empty() && r1.is_empty()
}
pub fn poll_fill_buf<T>(&mut self, cx: &mut Context<'_>) -> Poll<Region<T>>
where
S: Storage<T>,
{
if self.has_data() {
return Poll::Ready(self.buf());
}
self.state.data_available_waker.register(cx.waker());
if self.has_data() {
self.state.data_available_waker.take();
return Poll::Ready(self.buf());
}
if !self.is_writer_open() {
self.state.data_available_waker.take();
return Poll::Ready(Default::default());
}
Poll::Pending
}
pub fn consume(&mut self, amt: usize) {
self.state.ring.advance_left(amt);
self.state.space_available_waker.wake();
}
pub async fn read<T>(&mut self, buf: &mut [T]) -> usize
where
S: Storage<T>,
T: Clone,
{
Read { reader: self, buf }.await
}
pub async fn read_exact<T>(&mut self, buf: &mut [T]) -> Result<usize, ReadExactError>
where
S: Storage<T>,
T: Clone,
{
ReadExact {
reader: self,
buf,
read_bytes: 0,
}
.await
}
#[inline]
pub fn close(&mut self) {
self.state.close_reader();
}
#[inline]
fn buf<T>(&self) -> Region<T>
where
S: Storage<T>,
{
let (range_0, range_1) = self.state.ring.left_ranges();
Region::new(
self.state.storage.slice(range_0),
self.state.storage.slice(range_1),
)
}
}
impl<S> Drop for Reader<S> {
#[inline]
fn drop(&mut self) {
self.state.close_reader();
}
}
struct Read<'a, T, S> {
reader: &'a mut Reader<S>,
buf: &'a mut [T],
}
impl<'a, T, S> Future for Read<'a, T, S>
where
S: Storage<T>,
T: Clone,
{
type Output = usize;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.get_mut();
let src_buf = match me.reader.poll_fill_buf(cx) {
Poll::Ready(src_buf) => src_buf,
Poll::Pending => return Poll::Pending,
};
if src_buf.is_empty() {
return Poll::Ready(0);
}
let len = src_buf.len().min(me.buf.len());
src_buf.slice(..len).clone_to_slice(&mut me.buf[..len]);
me.reader.consume(len);
Poll::Ready(len)
}
}
struct ReadExact<'a, T, S> {
reader: &'a mut Reader<S>,
buf: &'a mut [T],
read_bytes: usize,
}
impl<'a, T, S> Future for ReadExact<'a, T, S>
where
S: Storage<T>,
T: Clone,
{
type Output = Result<usize, ReadExactError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.get_mut();
if me.buf.is_empty() {
debug_assert_eq!(me.read_bytes, 0);
return Poll::Ready(Ok(0));
}
let src_buf = match me.reader.poll_fill_buf(cx) {
Poll::Ready(src_buf) => src_buf,
Poll::Pending => return Poll::Pending,
};
let dest_buf = &mut me.buf[me.read_bytes..];
let read_len = dest_buf.len().min(src_buf.len());
src_buf
.slice(..read_len)
.clone_to_slice(&mut dest_buf[..read_len]);
me.reader.consume(read_len);
me.read_bytes += read_len;
if dest_buf.len() == read_len {
Poll::Ready(Ok(me.read_bytes))
} else if me.reader.is_writer_open() {
Poll::Pending
} else {
Poll::Ready(Err(ReadExactError::WriterClosed))
}
}
}
impl<S> Writer<S> {
#[inline]
pub fn is_reader_open(&self) -> bool {
self.state.is_reader_open.load(Ordering::Acquire)
}
#[inline]
pub fn has_space(&self) -> bool {
let (r0, r1) = self.state.ring.right_ranges();
!r0.is_empty() || !r1.is_empty()
}
#[inline]
pub fn is_flushed(&self) -> bool {
let (r0, r1) = self.state.ring.left_ranges();
r0.is_empty() && r1.is_empty()
}
fn get_flush_state(&self) -> Option<Result<(), WriteError>> {
if self.is_flushed() {
return Some(Ok(()));
}
if !self.is_reader_open() {
return Some(Err(WriteError::ReaderClosed));
}
None
}
#[inline]
fn buf<T>(&mut self) -> RegionMut<T>
where
S: Storage<T>,
{
let (range_0, range_1) = self.state.ring.right_ranges();
RegionMut::new(
unsafe { self.state.storage.slice_mut_unchecked(range_0) },
unsafe { self.state.storage.slice_mut_unchecked(range_1) },
)
}
pub fn poll_empty_buf<T>(&mut self, cx: &mut Context<'_>) -> Poll<RegionMut<T>>
where
S: Storage<T>,
{
if !self.is_reader_open() {
return Poll::Ready(Default::default());
}
if self.has_space() {
return Poll::Ready(self.buf());
}
self.state.space_available_waker.register(cx.waker());
if !self.is_reader_open() {
self.state.data_available_waker.take();
return Poll::Ready(Default::default());
}
if self.has_space() {
self.state.data_available_waker.take();
return Poll::Ready(self.buf());
}
Poll::Pending
}
pub fn feed(&mut self, len: usize) {
self.state.ring.advance_right(len);
self.state.data_available_waker.wake();
}
pub async fn write<T>(&mut self, buf: &[T]) -> usize
where
S: Storage<T>,
T: Clone,
{
Write { writer: self, buf }.await
}
pub async fn write_all<T>(&mut self, buf: &[T]) -> Result<usize, WriteError>
where
S: Storage<T>,
T: Clone,
{
WriteAll {
writer: self,
buf,
written_bytes: 0,
}
.await
}
pub fn poll_flush(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), WriteError>> {
if let Some(flush_state) = self.get_flush_state() {
return Poll::Ready(flush_state);
}
self.state.space_available_waker.register(cx.waker());
if let Some(flush_state) = self.get_flush_state() {
self.state.space_available_waker.take();
return Poll::Ready(flush_state);
}
Poll::Pending
}
pub async fn flush(&mut self) -> Result<(), WriteError> {
Flush(self).await
}
pub fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), WriteError>> {
match self.poll_flush(cx) {
Poll::Ready(Ok(())) => {}
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
Poll::Pending => return Poll::Pending,
}
self.state.close_writer();
Poll::Ready(Ok(()))
}
pub async fn close(&mut self) -> Result<(), WriteError> {
Flush(self).await
}
}
impl<S> Drop for Writer<S> {
#[inline]
fn drop(&mut self) {
self.state.close_writer();
}
}
struct Write<'a, T, S> {
writer: &'a mut Writer<S>,
buf: &'a [T],
}
impl<'a, T, S> Future for Write<'a, T, S>
where
S: Storage<T>,
T: Clone,
{
type Output = usize;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.get_mut();
let mut dest_buf = match me.writer.poll_empty_buf(cx) {
Poll::Ready(dest_buf) => dest_buf,
Poll::Pending => return Poll::Pending,
};
if dest_buf.is_empty() {
return Poll::Ready(0);
}
let len = dest_buf.len().min(me.buf.len());
dest_buf.slice_mut(..len).clone_from_slice(&me.buf[..len]);
me.writer.feed(len);
Poll::Ready(len)
}
}
struct WriteAll<'a, T, S> {
writer: &'a mut Writer<S>,
buf: &'a [T],
written_bytes: usize,
}
impl<'a, T, S> Future for WriteAll<'a, T, S>
where
S: Storage<T>,
T: Clone,
{
type Output = Result<usize, WriteError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let me = self.get_mut();
if me.buf.is_empty() {
debug_assert_eq!(me.written_bytes, 0);
return Poll::Ready(Ok(0));
}
let mut dest_buf = match me.writer.poll_empty_buf(cx) {
Poll::Ready(dest_buf) => dest_buf,
Poll::Pending => return Poll::Pending,
};
let src_buf = &me.buf[me.written_bytes..];
if dest_buf.is_empty() {
return Poll::Ready(Err(WriteError::ReaderClosed));
}
let write_len = dest_buf.len().min(src_buf.len());
dest_buf
.slice_mut(..write_len)
.clone_from_slice(&src_buf[..write_len]);
me.writer.feed(write_len);
me.written_bytes += write_len;
if src_buf.len() == write_len {
Poll::Ready(Ok(me.written_bytes))
} else {
Poll::Pending
}
}
}
struct Flush<'a, S>(&'a mut Writer<S>);
impl<'a, S> Future for Flush<'a, S> {
type Output = Result<(), WriteError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.get_mut().0.poll_flush(cx)
}
}
struct Close<'a, S>(&'a mut Writer<S>);
impl<'a, S> Future for Close<'a, S> {
type Output = Result<(), WriteError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.get_mut().0.poll_close(cx)
}
}
impl fmt::Display for WriteError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
WriteError::ReaderClosed => write!(f, "reader closed"),
}
}
}
impl fmt::Display for ReadExactError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ReadExactError::WriterClosed => write!(f, "writer closed"),
}
}
}
#[cfg(feature = "std")]
mod std_impls {
use crate::asyncio::{ReadExactError, WriteError};
use std::{error, io};
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl error::Error for WriteError {}
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl From<WriteError> for io::Error {
fn from(err: WriteError) -> Self {
match err {
WriteError::ReaderClosed => io::ErrorKind::UnexpectedEof.into(),
}
}
}
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl error::Error for ReadExactError {}
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl From<ReadExactError> for io::Error {
fn from(err: ReadExactError) -> Self {
match err {
ReadExactError::WriterClosed => io::ErrorKind::UnexpectedEof.into(),
}
}
}
}
#[cfg(feature = "std")]
pub use self::std_impls::*;
#[cfg(feature = "std-io")]
mod io_impls {
use crate::asyncio::{Reader, Writer};
use crate::storage::Storage;
use core::pin::Pin;
use core::task::{Context, Poll};
use futures::{io, AsyncBufRead, AsyncRead, AsyncWrite};
#[cfg_attr(docsrs, doc(cfg(feature = "std-io")))]
impl<S> AsyncRead for Reader<S>
where
S: Storage<u8>,
{
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
let me = self.get_mut();
let src_buf = match me.poll_fill_buf(cx) {
Poll::Ready(src_buf) => src_buf,
Poll::Pending => return Poll::Pending,
};
if src_buf.is_empty() {
return Poll::Ready(Ok(0));
}
let len = src_buf.len().min(buf.len());
src_buf.slice(..len).copy_to_slice(&mut buf[..len]);
me.consume(len);
Poll::Ready(Ok(len))
}
}
#[cfg_attr(docsrs, doc(cfg(feature = "std-io")))]
impl<S> AsyncBufRead for Reader<S>
where
S: Storage<u8>,
{
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
self.get_mut()
.poll_fill_buf(cx)
.map(|region| Ok(region.contiguous()))
}
fn consume(self: Pin<&mut Self>, amt: usize) {
self.get_mut().consume(amt);
}
}
#[cfg_attr(docsrs, doc(cfg(feature = "std-io")))]
impl<S> AsyncWrite for Writer<S>
where
S: Storage<u8>,
{
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
let me = self.get_mut();
let mut dest_buf = match me.poll_empty_buf(cx) {
Poll::Ready(dest_buf) => dest_buf,
Poll::Pending => return Poll::Pending,
};
if dest_buf.is_empty() {
return Poll::Ready(Ok(0));
}
let len = dest_buf.len().min(buf.len());
dest_buf.slice_mut(..len).copy_from_slice(&buf[..len]);
me.feed(len);
Poll::Ready(Ok(len))
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.get_mut().poll_flush(cx).map(|r| r.map_err(Into::into))
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.get_mut().poll_close(cx).map(|r| r.map_err(Into::into))
}
}
}
#[cfg(feature = "std-io")]
pub use self::io_impls::*;