use core::future::Future;
use core::marker::PhantomData;
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::task::{Context, Poll};
use super::{PubSubBehavior, PubSubChannel, WaitResult};
use crate::blocking_mutex::raw::RawMutex;
pub struct Sub<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> {
next_message_id: u64,
channel: &'a PSB,
_phantom: PhantomData<T>,
}
impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Sub<'a, PSB, T> {
pub(super) fn new(next_message_id: u64, channel: &'a PSB) -> Self {
Self {
next_message_id,
channel,
_phantom: Default::default(),
}
}
pub fn next_message<'s>(&'s mut self) -> SubscriberWaitFuture<'s, 'a, PSB, T> {
SubscriberWaitFuture { subscriber: self }
}
pub async fn next_message_pure(&mut self) -> T {
loop {
match self.next_message().await {
WaitResult::Lagged(_) => continue,
WaitResult::Message(message) => break message,
}
}
}
pub fn try_next_message(&mut self) -> Option<WaitResult<T>> {
match self.channel.get_message_with_context(&mut self.next_message_id, None) {
Poll::Ready(result) => Some(result),
Poll::Pending => None,
}
}
pub fn try_next_message_pure(&mut self) -> Option<T> {
loop {
match self.try_next_message() {
Some(WaitResult::Lagged(_)) => continue,
Some(WaitResult::Message(message)) => break Some(message),
None => break None,
}
}
}
pub fn available(&self) -> u64 {
self.channel.available(self.next_message_id)
}
}
impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Drop for Sub<'a, PSB, T> {
fn drop(&mut self) {
self.channel.unregister_subscriber(self.next_message_id)
}
}
impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Unpin for Sub<'a, PSB, T> {}
impl<'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> futures_util::Stream for Sub<'a, PSB, T> {
type Item = T;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match self
.channel
.get_message_with_context(&mut self.next_message_id, Some(cx))
{
Poll::Ready(WaitResult::Message(message)) => Poll::Ready(Some(message)),
Poll::Ready(WaitResult::Lagged(_)) => {
cx.waker().wake_by_ref();
Poll::Pending
}
Poll::Pending => Poll::Pending,
}
}
}
pub struct DynSubscriber<'a, T: Clone>(pub(super) Sub<'a, dyn PubSubBehavior<T> + 'a, T>);
impl<'a, T: Clone> Deref for DynSubscriber<'a, T> {
type Target = Sub<'a, dyn PubSubBehavior<T> + 'a, T>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<'a, T: Clone> DerefMut for DynSubscriber<'a, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
pub struct Subscriber<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize>(
pub(super) Sub<'a, PubSubChannel<M, T, CAP, SUBS, PUBS>, T>,
);
impl<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> Deref
for Subscriber<'a, M, T, CAP, SUBS, PUBS>
{
type Target = Sub<'a, PubSubChannel<M, T, CAP, SUBS, PUBS>, T>;
fn deref(&self) -> &Self::Target {
&self.0
}
}
impl<'a, M: RawMutex, T: Clone, const CAP: usize, const SUBS: usize, const PUBS: usize> DerefMut
for Subscriber<'a, M, T, CAP, SUBS, PUBS>
{
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.0
}
}
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct SubscriberWaitFuture<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> {
subscriber: &'s mut Sub<'a, PSB, T>,
}
impl<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Future for SubscriberWaitFuture<'s, 'a, PSB, T> {
type Output = WaitResult<T>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.subscriber
.channel
.get_message_with_context(&mut self.subscriber.next_message_id, Some(cx))
}
}
impl<'s, 'a, PSB: PubSubBehavior<T> + ?Sized, T: Clone> Unpin for SubscriberWaitFuture<'s, 'a, PSB, T> {}