#![allow(unsafe_code)]
use core::fmt;
use core::future::Future;
use core::mem::MaybeUninit;
use core::pin::Pin;
use core::task::{Context, Poll};
use super::atomic_waker::AtomicWaker;
use super::internal::{AtomicUsize, BlockArc, Ordering, UnsafeCell};
const VALUE_SENT: usize = 0b001;
const RX_CLOSED: usize = 0b010;
const TX_DROPPED: usize = 0b100;
pub mod error {
use core::fmt;
#[derive(Debug, Eq, PartialEq, Clone)]
pub struct RecvError(pub(super) ());
#[derive(Debug, Eq, PartialEq, Clone)]
pub enum TryRecvError {
Empty,
Closed,
}
impl fmt::Display for RecvError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "channel closed")
}
}
impl std::error::Error for RecvError {}
impl fmt::Display for TryRecvError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
TryRecvError::Empty => write!(fmt, "channel empty"),
TryRecvError::Closed => write!(fmt, "channel closed"),
}
}
}
impl std::error::Error for TryRecvError {}
}
use self::error::{RecvError, TryRecvError};
#[derive(Clone, Copy, PartialEq, Eq)]
enum RxClose {
Open,
WithValue,
Empty,
}
struct Inner<T> {
state: AtomicUsize,
value: UnsafeCell<MaybeUninit<T>>,
rx_waker: AtomicWaker,
tx_waker: AtomicWaker,
}
unsafe impl<T: Send> Send for Inner<T> {}
unsafe impl<T: Send> Sync for Inner<T> {}
impl<T> Inner<T> {
fn new() -> Self {
Self {
state: AtomicUsize::new(0),
value: UnsafeCell::new(MaybeUninit::uninit()),
rx_waker: AtomicWaker::new(),
tx_waker: AtomicWaker::new(),
}
}
unsafe fn read_value(&self) -> T {
self.value
.with_mut(|slot| unsafe { (*slot).assume_init_read() })
}
unsafe fn resolve(&self, close: RxClose, state: usize) -> Option<Option<T>> {
if close == RxClose::Empty {
Some(None)
} else if state & VALUE_SENT != 0 {
Some(Some(unsafe { self.read_value() }))
} else if state & TX_DROPPED != 0 {
Some(None)
} else {
None
}
}
}
pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
let inner = BlockArc::new(Inner::new());
let tx = Sender {
inner: Some(BlockArc::clone(&inner)),
};
let rx = Receiver {
inner: Some(inner),
close: RxClose::Open,
};
(tx, rx)
}
pub struct Sender<T> {
inner: Option<BlockArc<Inner<T>>>,
}
impl<T> Sender<T> {
pub fn send(mut self, value: T) -> Result<(), T> {
let Some(inner) = self.inner.take() else {
return Err(value);
};
inner.value.with_mut(|slot| unsafe {
(*slot).write(value);
});
let prev = inner
.state
.fetch_or(VALUE_SENT | TX_DROPPED, Ordering::AcqRel);
if prev & RX_CLOSED != 0 {
return Err(unsafe { inner.read_value() });
}
inner.rx_waker.wake();
Ok(())
}
pub async fn closed(&mut self) {
std::future::poll_fn(|cx| self.poll_closed(cx)).await;
}
pub fn is_closed(&self) -> bool {
self.inner
.as_ref()
.is_none_or(|inner| inner.state.load(Ordering::Acquire) & RX_CLOSED != 0)
}
pub fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<()> {
let Some(inner) = self.inner.as_ref() else {
return Poll::Ready(());
};
if inner.state.load(Ordering::Acquire) & RX_CLOSED != 0 {
return Poll::Ready(());
}
inner.tx_waker.register(cx.waker());
if inner.state.load(Ordering::Acquire) & RX_CLOSED != 0 {
Poll::Ready(())
} else {
Poll::Pending
}
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
let Some(inner) = self.inner.as_ref() else {
return;
};
let prev = inner.state.fetch_or(TX_DROPPED, Ordering::AcqRel);
if prev & RX_CLOSED == 0 {
inner.rx_waker.wake();
}
}
}
impl<T> fmt::Debug for Sender<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Sender")
.field("is_closed", &self.is_closed())
.finish()
}
}
pub struct Receiver<T> {
inner: Option<BlockArc<Inner<T>>>,
close: RxClose,
}
impl<T> Receiver<T> {
pub fn close(&mut self) {
let Some(inner) = self.inner.as_ref() else {
return;
};
if self.close != RxClose::Open {
return;
}
let prev = inner.state.fetch_or(RX_CLOSED, Ordering::AcqRel);
self.close = if prev & VALUE_SENT != 0 {
RxClose::WithValue
} else {
RxClose::Empty
};
if prev & TX_DROPPED == 0 {
inner.tx_waker.wake();
}
}
pub fn is_terminated(&self) -> bool {
self.inner.is_none()
}
pub fn is_empty(&self) -> bool {
let Some(inner) = self.inner.as_ref() else {
return true;
};
self.close == RxClose::Empty || inner.state.load(Ordering::Acquire) & VALUE_SENT == 0
}
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
let Some(inner) = self.inner.as_ref() else {
return Err(TryRecvError::Closed);
};
let state = inner.state.load(Ordering::Acquire);
match unsafe { inner.resolve(self.close, state) } {
Some(done) => {
self.inner = None;
done.ok_or(TryRecvError::Closed)
}
None => Err(TryRecvError::Empty),
}
}
pub fn blocking_recv(self) -> Result<T, RecvError> {
super::internal::blocking(self)
}
}
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
let Some(inner) = self.inner.as_ref() else {
return;
};
let owns_value = match self.close {
RxClose::WithValue => true,
RxClose::Empty => false,
RxClose::Open => {
let prev = inner.state.fetch_or(RX_CLOSED, Ordering::AcqRel);
if prev & TX_DROPPED == 0 {
inner.tx_waker.wake();
}
prev & VALUE_SENT != 0
}
};
if owns_value {
drop(unsafe { inner.read_value() });
}
}
}
impl<T> Future for Receiver<T> {
type Output = Result<T, RecvError>;
#[expect(
clippy::panic,
reason = "polled after it completed, which the `Future` contract rules out"
)]
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
let Some(inner) = this.inner.as_ref() else {
panic!("called after complete");
};
let mut state = inner.state.load(Ordering::Acquire);
let mut registered = false;
loop {
if let Some(done) = unsafe { inner.resolve(this.close, state) } {
this.inner = None;
return Poll::Ready(done.ok_or(RecvError(())));
}
if registered {
return Poll::Pending;
}
inner.rx_waker.register(cx.waker());
registered = true;
state = inner.state.load(Ordering::Acquire);
}
}
}
impl<T> fmt::Debug for Receiver<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Receiver")
.field("is_terminated", &self.is_terminated())
.field("is_empty", &self.is_empty())
.finish()
}
}
#[cfg(all(test, not(loom)))]
mod tests {
use super::*;
use crate::sync::test_support::{Polled, block_on};
use std::sync::atomic::{AtomicUsize as StdAtomicUsize, Ordering as StdOrdering};
struct DropCount(std::sync::Arc<StdAtomicUsize>);
impl Drop for DropCount {
fn drop(&mut self) {
self.0.fetch_add(1, StdOrdering::SeqCst);
}
}
fn counted() -> (std::sync::Arc<StdAtomicUsize>, DropCount) {
let drops = std::sync::Arc::new(StdAtomicUsize::new(0));
(std::sync::Arc::clone(&drops), DropCount(drops))
}
#[test]
fn a_pending_receiver_is_woken_by_the_send() {
let (tx, rx) = channel();
let mut rx = Polled::new(rx);
rx.pending();
assert!(tx.send(5).is_ok());
assert_eq!(rx.wakes(), 1);
assert_eq!(rx.ready(), Ok(5));
}
#[test]
fn a_pending_sender_is_woken_by_the_close() {
let (mut tx, mut rx) = channel::<u8>();
let mut closed = Polled::new(std::future::poll_fn(|cx| tx.poll_closed(cx)));
closed.pending();
rx.close();
assert_eq!(closed.wakes(), 1);
closed.ready();
drop(closed);
assert_eq!(tx.send(1), Err(1));
}
#[test]
fn an_unreceived_value_is_dropped_by_the_receiver() {
let (drops, value) = counted();
let (tx, rx) = channel();
assert!(tx.send(value).is_ok());
assert_eq!(drops.load(StdOrdering::SeqCst), 0, "the channel holds it");
drop(rx);
assert_eq!(drops.load(StdOrdering::SeqCst), 1);
}
#[test]
fn a_closed_receiver_still_owns_a_value_sent_before_the_close() {
let (drops, value) = counted();
let (tx, mut rx) = channel();
assert!(tx.send(value).is_ok());
rx.close();
assert_eq!(drops.load(StdOrdering::SeqCst), 0);
drop(rx);
assert_eq!(drops.load(StdOrdering::SeqCst), 1);
}
#[test]
fn a_failed_send_returns_the_value_and_drops_it_once() {
let (drops, value) = counted();
let (tx, mut rx) = channel();
rx.close();
let back = tx.send(value).expect_err("the receiver closed");
assert_eq!(drops.load(StdOrdering::SeqCst), 0, "returned, not dropped");
drop(back);
drop(rx);
assert_eq!(drops.load(StdOrdering::SeqCst), 1);
}
#[test]
fn a_received_value_is_not_dropped_by_the_channel() {
let (drops, value) = counted();
let (tx, mut rx) = channel();
assert!(tx.send(value).is_ok());
let value = rx.try_recv().expect("sent");
drop(rx);
assert_eq!(drops.load(StdOrdering::SeqCst), 0);
drop(value);
assert_eq!(drops.load(StdOrdering::SeqCst), 1);
}
#[test]
fn a_value_sent_across_threads_arrives_and_drops_once() {
for _ in 0..if cfg!(miri) { 5 } else { 500 } {
let (drops, value) = counted();
let (tx, rx) = channel();
let sender = std::thread::spawn(move || tx.send(value).is_ok());
let got = block_on(rx);
assert!(sender.join().expect("sender"));
drop(got);
assert_eq!(drops.load(StdOrdering::SeqCst), 1);
}
}
#[test]
#[should_panic(expected = "called after complete")]
fn polling_a_completed_receiver_panics() {
let (tx, rx) = channel();
assert!(tx.send(1).is_ok());
let mut rx = Polled::new(rx);
assert_eq!(rx.ready(), Ok(1));
let _ = rx.poll();
}
}