pub mod oneshot {
use std::{
cell::RefCell,
future::Future,
pin::Pin,
rc::Rc,
task::{Context, Poll, Waker},
};
struct OneshotState<T> {
value: Option<T>,
waker: Option<Waker>,
sender_dropped: bool,
receiver_dropped: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct RecvError;
impl std::fmt::Display for RecvError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "oneshot sender dropped without sending a value")
}
}
impl std::error::Error for RecvError {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TryRecvError {
Empty,
Closed,
}
impl std::fmt::Display for TryRecvError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Empty => write!(f, "oneshot channel is empty"),
Self::Closed => write!(f, "oneshot sender was dropped"),
}
}
}
impl std::error::Error for TryRecvError {}
pub struct Sender<T> {
state: Rc<RefCell<OneshotState<T>>>,
}
impl<T> std::fmt::Debug for Sender<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Sender")
.field("is_closed", &self.is_closed())
.finish()
}
}
impl<T> Sender<T> {
pub fn send(self, value: T) -> Result<(), T> {
let mut state = self.state.borrow_mut();
if state.receiver_dropped {
return Err(value);
}
state.value = Some(value);
if let Some(w) = state.waker.take() {
w.wake();
}
Ok(())
}
pub fn is_closed(&self) -> bool {
self.state.borrow().receiver_dropped
}
pub fn is_canceled(&self) -> bool {
self.is_closed()
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
let mut state = self.state.borrow_mut();
state.sender_dropped = true;
if state.value.is_none() {
if let Some(w) = state.waker.take() {
w.wake();
}
}
}
}
pub struct Receiver<T> {
state: Rc<RefCell<OneshotState<T>>>,
}
impl<T> std::fmt::Debug for Receiver<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Receiver").finish()
}
}
impl<T> Receiver<T> {
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
let mut state = self.state.borrow_mut();
if let Some(val) = state.value.take() {
Ok(val)
} else if state.sender_dropped {
Err(TryRecvError::Closed)
} else {
Err(TryRecvError::Empty)
}
}
pub fn close(&mut self) {
let mut state = self.state.borrow_mut();
state.receiver_dropped = true;
state.waker = None;
state.value = None;
}
}
impl<T> Future for Receiver<T> {
type Output = Result<T, RecvError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut state = self.state.borrow_mut();
if let Some(val) = state.value.take() {
Poll::Ready(Ok(val))
} else if state.sender_dropped {
Poll::Ready(Err(RecvError))
} else {
state.waker = Some(cx.waker().clone());
Poll::Pending
}
}
}
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
let mut state = self.state.borrow_mut();
state.receiver_dropped = true;
state.waker = None;
state.value = None;
}
}
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
let state = Rc::new(RefCell::new(OneshotState {
value: None,
waker: None,
sender_dropped: false,
receiver_dropped: false,
}));
(
Sender {
state: Rc::clone(&state),
},
Receiver { state },
)
}
}
pub mod mpsc {
use std::{
cell::RefCell,
collections::VecDeque,
future::Future,
pin::Pin,
rc::Rc,
task::{Context, Poll, Waker},
};
struct ChannelState<T> {
queue: VecDeque<T>,
capacity: usize,
receiver_waker: Option<Waker>,
sender_wakers: std::collections::BTreeMap<u64, Waker>,
next_waiter_id: u64,
sender_count: usize,
receiver_count: usize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SendError<T>(pub T);
impl<T> std::fmt::Display for SendError<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "channel closed")
}
}
impl<T: std::fmt::Debug> std::error::Error for SendError<T> {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TrySendError<T> {
Full(T),
Closed(T),
}
impl<T> std::fmt::Display for TrySendError<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Full(_) => write!(f, "channel buffer is full"),
Self::Closed(_) => write!(f, "channel closed"),
}
}
}
impl<T: std::fmt::Debug> std::error::Error for TrySendError<T> {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TryRecvError {
Empty,
Disconnected,
}
impl std::fmt::Display for TryRecvError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Empty => write!(f, "channel is empty"),
Self::Disconnected => write!(f, "channel is disconnected"),
}
}
}
impl std::error::Error for TryRecvError {}
pub struct Sender<T> {
state: Rc<RefCell<ChannelState<T>>>,
}
impl<T> std::fmt::Debug for Sender<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Sender")
.field("capacity", &self.capacity())
.field("is_closed", &self.is_closed())
.finish()
}
}
impl<T> Clone for Sender<T> {
fn clone(&self) -> Self {
self.state.borrow_mut().sender_count += 1;
Self {
state: Rc::clone(&self.state),
}
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
let mut state = self.state.borrow_mut();
state.sender_count = state.sender_count.saturating_sub(1);
if state.sender_count == 0 {
if let Some(w) = state.receiver_waker.take() {
w.wake();
}
}
}
}
impl<T> Sender<T> {
pub fn send(&self, value: T) -> Send<'_, T> {
Send {
sender: self,
value: Some(value),
waiter_id: None,
}
}
pub fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
let mut state = self.state.borrow_mut();
if state.receiver_count == 0 {
return Err(TrySendError::Closed(value));
}
if state.queue.len() >= state.capacity || !state.sender_wakers.is_empty() {
return Err(TrySendError::Full(value));
}
state.queue.push_back(value);
if let Some(w) = state.receiver_waker.take() {
w.wake();
}
Ok(())
}
pub fn is_closed(&self) -> bool {
self.state.borrow().receiver_count == 0
}
pub fn capacity(&self) -> usize {
self.state.borrow().capacity
}
}
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Send<'a, T> {
sender: &'a Sender<T>,
value: Option<T>,
waiter_id: Option<u64>,
}
impl<'a, T> std::fmt::Debug for Send<'a, T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Send")
.field("has_value", &self.value.is_some())
.finish()
}
}
impl<'a, T: Unpin> Future for Send<'a, T> {
type Output = Result<(), SendError<T>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
let mut state = this.sender.state.borrow_mut();
if state.receiver_count == 0 {
let val = this.value.take().expect("polled after completion");
return Poll::Ready(Err(SendError(val)));
}
if state.queue.len() < state.capacity {
let val = this.value.take().expect("polled after completion");
state.queue.push_back(val);
if let Some(w) = state.receiver_waker.take() {
w.wake();
}
return Poll::Ready(Ok(()));
}
let waiter_id = match this.waiter_id {
Some(id) => id,
None => {
let id = state.next_waiter_id;
state.next_waiter_id = state.next_waiter_id.wrapping_add(1);
this.waiter_id = Some(id);
id
}
};
match state.sender_wakers.get_mut(&waiter_id) {
Some(w) => {
if !w.will_wake(cx.waker()) {
*w = cx.waker().clone();
}
}
None => {
state.sender_wakers.insert(waiter_id, cx.waker().clone());
}
}
Poll::Pending
}
}
impl<'a, T> Drop for Send<'a, T> {
fn drop(&mut self) {
if let Some(id) = self.waiter_id {
let mut state = self.sender.state.borrow_mut();
state.sender_wakers.remove(&id);
}
}
}
pub struct Receiver<T> {
state: Rc<RefCell<ChannelState<T>>>,
}
impl<T> std::fmt::Debug for Receiver<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Receiver").finish()
}
}
impl<T> Receiver<T> {
pub async fn recv(&mut self) -> Option<T> {
struct Recv<'a, T> {
receiver: &'a mut Receiver<T>,
}
impl<'a, T> Future for Recv<'a, T> {
type Output = Option<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.receiver.poll_recv(cx)
}
}
impl<'a, T> Drop for Recv<'a, T> {
fn drop(&mut self) {
let mut state = self.receiver.state.borrow_mut();
state.receiver_waker = None;
}
}
Recv { receiver: self }.await
}
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
let mut state = self.state.borrow_mut();
if let Some(val) = state.queue.pop_front() {
if let Some((_id, waker)) = state.sender_wakers.pop_first() {
waker.wake();
}
Poll::Ready(Some(val))
} else if state.sender_count == 0 {
Poll::Ready(None)
} else {
state.receiver_waker = Some(cx.waker().clone());
Poll::Pending
}
}
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
let mut state = self.state.borrow_mut();
if let Some(val) = state.queue.pop_front() {
if let Some((_id, waker)) = state.sender_wakers.pop_first() {
waker.wake();
}
Ok(val)
} else if state.sender_count == 0 {
Err(TryRecvError::Disconnected)
} else {
Err(TryRecvError::Empty)
}
}
pub fn close(&mut self) {
let mut state = self.state.borrow_mut();
state.receiver_count = 0;
state.receiver_waker = None;
state.queue.clear();
while let Some((_id, waker)) = state.sender_wakers.pop_first() {
waker.wake();
}
}
}
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
let mut state = self.state.borrow_mut();
state.receiver_count = 0;
state.receiver_waker = None;
state.queue.clear();
while let Some((_id, waker)) = state.sender_wakers.pop_first() {
waker.wake();
}
}
}
pub fn channel<T>(buffer: usize) -> (Sender<T>, Receiver<T>) {
assert!(buffer > 0, "mpsc channel capacity must be greater than 0");
let state = Rc::new(RefCell::new(ChannelState {
queue: VecDeque::with_capacity(buffer),
capacity: buffer,
receiver_waker: None,
sender_wakers: std::collections::BTreeMap::new(),
next_waiter_id: 0,
sender_count: 1,
receiver_count: 1,
}));
(
Sender {
state: Rc::clone(&state),
},
Receiver { state },
)
}
pub struct UnboundedSender<T> {
state: Rc<RefCell<ChannelState<T>>>,
}
impl<T> std::fmt::Debug for UnboundedSender<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("UnboundedSender")
.field("is_closed", &self.is_closed())
.finish()
}
}
impl<T> Clone for UnboundedSender<T> {
fn clone(&self) -> Self {
self.state.borrow_mut().sender_count += 1;
Self {
state: Rc::clone(&self.state),
}
}
}
impl<T> Drop for UnboundedSender<T> {
fn drop(&mut self) {
let mut state = self.state.borrow_mut();
state.sender_count = state.sender_count.saturating_sub(1);
if state.sender_count == 0 {
if let Some(w) = state.receiver_waker.take() {
w.wake();
}
}
}
}
impl<T> UnboundedSender<T> {
pub fn send(&self, value: T) -> Result<(), SendError<T>> {
let mut state = self.state.borrow_mut();
if state.receiver_count == 0 {
return Err(SendError(value));
}
state.queue.push_back(value);
if let Some(w) = state.receiver_waker.take() {
w.wake();
}
Ok(())
}
pub fn is_closed(&self) -> bool {
self.state.borrow().receiver_count == 0
}
}
pub struct UnboundedReceiver<T> {
state: Rc<RefCell<ChannelState<T>>>,
}
impl<T> std::fmt::Debug for UnboundedReceiver<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("UnboundedReceiver").finish()
}
}
impl<T> UnboundedReceiver<T> {
pub async fn recv(&mut self) -> Option<T> {
struct UnboundedRecv<'a, T> {
receiver: &'a mut UnboundedReceiver<T>,
}
impl<'a, T> Future for UnboundedRecv<'a, T> {
type Output = Option<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.receiver.poll_recv(cx)
}
}
impl<'a, T> Drop for UnboundedRecv<'a, T> {
fn drop(&mut self) {
let mut state = self.receiver.state.borrow_mut();
state.receiver_waker = None;
}
}
UnboundedRecv { receiver: self }.await
}
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
let mut state = self.state.borrow_mut();
if let Some(val) = state.queue.pop_front() {
Poll::Ready(Some(val))
} else if state.sender_count == 0 {
Poll::Ready(None)
} else {
state.receiver_waker = Some(cx.waker().clone());
Poll::Pending
}
}
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
let mut state = self.state.borrow_mut();
if let Some(val) = state.queue.pop_front() {
Ok(val)
} else if state.sender_count == 0 {
Err(TryRecvError::Disconnected)
} else {
Err(TryRecvError::Empty)
}
}
pub fn close(&mut self) {
let mut state = self.state.borrow_mut();
state.receiver_count = 0;
state.receiver_waker = None;
state.queue.clear();
}
}
impl<T> Drop for UnboundedReceiver<T> {
fn drop(&mut self) {
let mut state = self.state.borrow_mut();
state.receiver_count = 0;
state.receiver_waker = None;
state.queue.clear();
}
}
pub fn unbounded_channel<T>() -> (UnboundedSender<T>, UnboundedReceiver<T>) {
let state = Rc::new(RefCell::new(ChannelState {
queue: VecDeque::new(),
capacity: usize::MAX,
receiver_waker: None,
sender_wakers: std::collections::BTreeMap::new(),
next_waiter_id: 0,
sender_count: 1,
receiver_count: 1,
}));
(
UnboundedSender {
state: Rc::clone(&state),
},
UnboundedReceiver { state },
)
}
}