use super::{ChannelSendError, CloseStatus};
use crate::{
intrusive_double_linked_list::{LinkedList, ListNode},
utils::update_waker_ref,
NoopLock,
};
use core::marker::PhantomData;
use core::pin::Pin;
use futures_core::{
future::{FusedFuture, Future},
task::{Context, Poll, Waker},
};
use lock_api::{Mutex, RawMutex};
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Ord, PartialOrd)]
pub struct StateId(u64);
impl StateId {
pub fn new() -> Self {
StateId(0)
}
}
#[derive(PartialEq, Debug)]
pub enum RecvPollState {
Unregistered,
Registered,
}
#[derive(Debug)]
pub struct RecvWaitQueueEntry {
task: Option<Waker>,
state: RecvPollState,
state_id: StateId,
}
impl RecvWaitQueueEntry {
pub fn new(state_id: StateId) -> RecvWaitQueueEntry {
RecvWaitQueueEntry {
task: None,
state_id,
state: RecvPollState::Unregistered,
}
}
}
pub trait ChannelReceiveAccess<T> {
unsafe fn receive_or_register(
&self,
wait_node: &mut ListNode<RecvWaitQueueEntry>,
cx: &mut Context<'_>,
) -> Poll<Option<(StateId, T)>>;
fn remove_receive_waiter(
&self,
wait_node: &mut ListNode<RecvWaitQueueEntry>,
);
}
#[must_use = "futures do nothing unless polled"]
pub struct StateReceiveFuture<'a, MutexType, T>
where
T: Clone,
{
channel: Option<&'a dyn ChannelReceiveAccess<T>>,
wait_node: ListNode<RecvWaitQueueEntry>,
_phantom: PhantomData<MutexType>,
}
unsafe impl<'a, MutexType: Sync, T: Clone + Send> Send
for StateReceiveFuture<'a, MutexType, T>
{
}
impl<'a, MutexType, T: Clone> core::fmt::Debug
for StateReceiveFuture<'a, MutexType, T>
{
fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
f.debug_struct("StateReceiveFuture").finish()
}
}
impl<'a, MutexType, T: Clone> Future for StateReceiveFuture<'a, MutexType, T> {
type Output = Option<(StateId, T)>;
fn poll(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<(StateId, T)>> {
let mut_self: &mut StateReceiveFuture<MutexType, T> =
unsafe { Pin::get_unchecked_mut(self) };
let channel = mut_self
.channel
.expect("polled StateReceiveFuture after completion");
let poll_res =
unsafe { channel.receive_or_register(&mut mut_self.wait_node, cx) };
if poll_res.is_ready() {
mut_self.channel = None;
}
poll_res
}
}
impl<'a, MutexType, T: Clone> FusedFuture
for StateReceiveFuture<'a, MutexType, T>
{
fn is_terminated(&self) -> bool {
self.channel.is_none()
}
}
impl<'a, MutexType, T: Clone> Drop for StateReceiveFuture<'a, MutexType, T> {
fn drop(&mut self) {
if let Some(channel) = self.channel {
channel.remove_receive_waiter(&mut self.wait_node);
}
}
}
fn wake_waiters(waiters: &mut LinkedList<RecvWaitQueueEntry>) {
waiters.reverse_drain(|waiter| {
if let Some(handle) = waiter.task.take() {
handle.wake();
}
waiter.state = RecvPollState::Unregistered;
});
}
struct ChannelState<T> {
is_closed: bool,
state_id: StateId,
value: Option<T>,
waiters: LinkedList<RecvWaitQueueEntry>,
}
impl<T> ChannelState<T>
where
T: Clone,
{
fn new() -> ChannelState<T> {
ChannelState::<T> {
is_closed: false,
state_id: StateId(0),
value: None,
waiters: LinkedList::new(),
}
}
fn send(&mut self, value: T) -> Result<(), ChannelSendError<T>> {
if self.is_closed || self.state_id.0 == core::u64::MAX {
return Err(ChannelSendError(value));
}
self.value = Some(value);
self.state_id.0 += 1;
wake_waiters(&mut self.waiters);
Ok(())
}
fn close(&mut self) -> CloseStatus {
if self.is_closed {
return CloseStatus::AlreadyClosed;
}
self.is_closed = true;
wake_waiters(&mut self.waiters);
CloseStatus::NewlyClosed
}
fn try_receive(&mut self, state_id: StateId) -> Option<(StateId, T)> {
let val = self.value.as_ref()?;
if state_id < self.state_id {
Some((self.state_id, val.clone()))
} else {
None
}
}
unsafe fn receive_or_register(
&mut self,
wait_node: &mut ListNode<RecvWaitQueueEntry>,
cx: &mut Context<'_>,
) -> Poll<Option<(StateId, T)>> {
match wait_node.state {
RecvPollState::Unregistered => {
let val_to_deliver = match &self.value {
Some(ref v) if wait_node.state_id < self.state_id => {
Some(v.clone())
}
Some(_) | None => None,
};
match val_to_deliver {
Some(v) => {
Poll::Ready(Some((self.state_id, v)))
}
None => {
if self.is_closed {
Poll::Ready(None)
} else {
wait_node.task = Some(cx.waker().clone());
wait_node.state = RecvPollState::Registered;
self.waiters.add_front(wait_node);
Poll::Pending
}
}
}
}
RecvPollState::Registered => {
update_waker_ref(&mut wait_node.task, cx);
Poll::Pending
}
}
}
fn remove_waiter(&mut self, wait_node: &mut ListNode<RecvWaitQueueEntry>) {
if let RecvPollState::Registered = wait_node.state {
if !unsafe { self.waiters.remove(wait_node) } {
panic!("Future could not be removed from wait queue");
}
wait_node.state = RecvPollState::Unregistered;
}
}
}
pub struct GenericStateBroadcastChannel<MutexType: RawMutex, T> {
inner: Mutex<MutexType, ChannelState<T>>,
}
unsafe impl<MutexType: RawMutex + Send, T: Send> Send
for GenericStateBroadcastChannel<MutexType, T>
{
}
unsafe impl<MutexType: RawMutex + Sync, T: Send> Sync
for GenericStateBroadcastChannel<MutexType, T>
{
}
impl<MutexType: RawMutex, T> core::fmt::Debug
for GenericStateBroadcastChannel<MutexType, T>
{
fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
f.debug_struct("GenericStateBroadcastChannel").finish()
}
}
impl<MutexType: RawMutex, T> GenericStateBroadcastChannel<MutexType, T>
where
T: Clone,
{
pub fn new() -> GenericStateBroadcastChannel<MutexType, T>
where
T: Clone,
{
GenericStateBroadcastChannel {
inner: Mutex::new(ChannelState::new()),
}
}
pub fn send(&self, value: T) -> Result<(), ChannelSendError<T>> {
self.inner.lock().send(value)
}
pub fn close(&self) -> CloseStatus {
self.inner.lock().close()
}
pub fn receive(
&self,
state_id: StateId,
) -> StateReceiveFuture<MutexType, T> {
StateReceiveFuture {
channel: Some(self),
wait_node: ListNode::new(RecvWaitQueueEntry::new(state_id)),
_phantom: PhantomData,
}
}
pub fn try_receive(&self, state_id: StateId) -> Option<(StateId, T)> {
self.inner.lock().try_receive(state_id)
}
}
impl<MutexType: RawMutex, T: Clone> ChannelReceiveAccess<T>
for GenericStateBroadcastChannel<MutexType, T>
{
unsafe fn receive_or_register(
&self,
wait_node: &mut ListNode<RecvWaitQueueEntry>,
cx: &mut Context<'_>,
) -> Poll<Option<(StateId, T)>> {
self.inner.lock().receive_or_register(wait_node, cx)
}
fn remove_receive_waiter(
&self,
wait_node: &mut ListNode<RecvWaitQueueEntry>,
) {
self.inner.lock().remove_waiter(wait_node)
}
}
pub type LocalStateBroadcastChannel<T> =
GenericStateBroadcastChannel<NoopLock, T>;
#[cfg(feature = "std")]
mod if_std {
use super::*;
pub type StateBroadcastChannel<T> =
GenericStateBroadcastChannel<parking_lot::RawMutex, T>;
}
#[cfg(feature = "std")]
pub use self::if_std::*;
#[cfg(feature = "alloc")]
mod if_alloc {
use super::*;
pub mod shared {
use super::*;
use core::sync::atomic::{AtomicUsize, Ordering};
struct GenericStateBroadcastChannelSharedState<MutexType, T>
where
MutexType: RawMutex,
T: Clone + 'static,
{
senders: AtomicUsize,
receivers: AtomicUsize,
channel: GenericStateBroadcastChannel<MutexType, T>,
}
impl<MutexType, T> ChannelReceiveAccess<T>
for GenericStateBroadcastChannelSharedState<MutexType, T>
where
MutexType: RawMutex,
T: Clone + 'static,
{
unsafe fn receive_or_register(
&self,
wait_node: &mut ListNode<RecvWaitQueueEntry>,
cx: &mut Context<'_>,
) -> Poll<Option<(StateId, T)>> {
self.channel.receive_or_register(wait_node, cx)
}
fn remove_receive_waiter(
&self,
wait_node: &mut ListNode<RecvWaitQueueEntry>,
) {
self.channel.remove_receive_waiter(wait_node)
}
}
#[must_use = "futures do nothing unless polled"]
pub struct StateReceiveFuture<MutexType, T> {
channel: Option<alloc::sync::Arc<dyn ChannelReceiveAccess<T>>>,
wait_node: ListNode<RecvWaitQueueEntry>,
_phantom: PhantomData<MutexType>,
}
unsafe impl<MutexType: Sync, T: Clone + Send> Send
for StateReceiveFuture<MutexType, T>
{
}
impl<MutexType, T> core::fmt::Debug for StateReceiveFuture<MutexType, T> {
fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
f.debug_struct("StateReceiveFuture").finish()
}
}
impl<MutexType, T> Future for StateReceiveFuture<MutexType, T> {
type Output = Option<(StateId, T)>;
fn poll(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<(StateId, T)>> {
let mut_self: &mut StateReceiveFuture<MutexType, T> =
unsafe { Pin::get_unchecked_mut(self) };
let channel = mut_self
.channel
.take()
.expect("polled StateReceiveFuture after completion");
let poll_res = unsafe {
channel.receive_or_register(&mut mut_self.wait_node, cx)
};
if poll_res.is_ready() {
mut_self.channel = None;
} else {
mut_self.channel = Some(channel)
}
poll_res
}
}
impl<MutexType, T> FusedFuture for StateReceiveFuture<MutexType, T> {
fn is_terminated(&self) -> bool {
self.channel.is_none()
}
}
impl<MutexType, T> Drop for StateReceiveFuture<MutexType, T> {
fn drop(&mut self) {
if let Some(channel) = &self.channel {
channel.remove_receive_waiter(&mut self.wait_node);
}
}
}
pub struct GenericStateSender<MutexType, T>
where
MutexType: RawMutex,
T: Clone + 'static,
{
inner: alloc::sync::Arc<
GenericStateBroadcastChannelSharedState<MutexType, T>,
>,
}
pub struct GenericStateReceiver<MutexType, T>
where
MutexType: RawMutex,
T: Clone + 'static,
{
inner: alloc::sync::Arc<
GenericStateBroadcastChannelSharedState<MutexType, T>,
>,
}
impl<MutexType, T> core::fmt::Debug for GenericStateSender<MutexType, T>
where
MutexType: RawMutex,
T: Clone,
{
fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
f.debug_struct("StateSender").finish()
}
}
impl<MutexType, T> core::fmt::Debug for GenericStateReceiver<MutexType, T>
where
MutexType: RawMutex,
T: Clone,
{
fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
f.debug_struct("StateReceiver").finish()
}
}
impl<MutexType, T> Clone for GenericStateSender<MutexType, T>
where
MutexType: RawMutex,
T: Clone,
{
fn clone(&self) -> Self {
let old_size =
self.inner.senders.fetch_add(1, Ordering::Relaxed);
if old_size > (core::isize::MAX) as usize {
panic!("Reached maximum refcount");
}
GenericStateSender {
inner: self.inner.clone(),
}
}
}
impl<MutexType, T> Drop for GenericStateSender<MutexType, T>
where
MutexType: RawMutex,
T: Clone,
{
fn drop(&mut self) {
if self.inner.senders.fetch_sub(1, Ordering::Release) != 1 {
return;
}
core::sync::atomic::fence(Ordering::Acquire);
self.inner.channel.close();
}
}
impl<MutexType, T> Clone for GenericStateReceiver<MutexType, T>
where
MutexType: RawMutex,
T: Clone,
{
fn clone(&self) -> Self {
let old_size =
self.inner.receivers.fetch_add(1, Ordering::Relaxed);
if old_size > (core::isize::MAX) as usize {
panic!("Reached maximum refcount");
}
GenericStateReceiver {
inner: self.inner.clone(),
}
}
}
impl<MutexType, T> Drop for GenericStateReceiver<MutexType, T>
where
MutexType: RawMutex,
T: Clone,
{
fn drop(&mut self) {
if self.inner.receivers.fetch_sub(1, Ordering::Release) != 1 {
return;
}
core::sync::atomic::fence(Ordering::Acquire);
self.inner.channel.close();
}
}
pub fn generic_state_broadcast_channel<MutexType, T>() -> (
GenericStateSender<MutexType, T>,
GenericStateReceiver<MutexType, T>,
)
where
MutexType: RawMutex,
T: Clone + Send,
{
let inner = alloc::sync::Arc::new(
GenericStateBroadcastChannelSharedState {
channel: GenericStateBroadcastChannel::new(),
senders: AtomicUsize::new(1),
receivers: AtomicUsize::new(1),
},
);
let sender = GenericStateSender {
inner: inner.clone(),
};
let receiver = GenericStateReceiver { inner };
(sender, receiver)
}
impl<MutexType, T> GenericStateSender<MutexType, T>
where
MutexType: RawMutex + 'static,
T: Clone,
{
pub fn send(&self, value: T) -> Result<(), ChannelSendError<T>> {
self.inner.channel.send(value)
}
}
impl<MutexType, T> GenericStateReceiver<MutexType, T>
where
MutexType: RawMutex + 'static,
T: Clone,
{
pub fn receive(
&self,
state_id: StateId,
) -> StateReceiveFuture<MutexType, T> {
StateReceiveFuture {
channel: Some(self.inner.clone()),
wait_node: ListNode::new(RecvWaitQueueEntry::new(state_id)),
_phantom: PhantomData,
}
}
pub fn try_receive(
&self,
state_id: StateId,
) -> Option<(StateId, T)> {
self.inner.channel.try_receive(state_id)
}
}
#[cfg(feature = "std")]
mod if_std {
use super::*;
pub type StateSender<T> =
GenericStateSender<parking_lot::RawMutex, T>;
pub type StateReceiver<T> =
GenericStateReceiver<parking_lot::RawMutex, T>;
pub fn state_broadcast_channel<T>(
) -> (StateSender<T>, StateReceiver<T>)
where
T: Clone + Send,
{
generic_state_broadcast_channel::<parking_lot::RawMutex, T>()
}
}
#[cfg(feature = "std")]
pub use self::if_std::*;
}
}
#[cfg(feature = "alloc")]
pub use self::if_alloc::*;