use core::fmt;
use crate::prelude::*;
pub type BoxNotify<'a, T = ()> = Pin<Box<dyn Notify<Event = T> + Send + 'a>>;
impl<T> fmt::Debug for BoxNotify<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("BoxNotify")
}
}
pub type LocalBoxNotify<'a, T = ()> = Pin<Box<dyn Notify<Event = T> + 'a>>;
impl<T> fmt::Debug for LocalBoxNotify<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("LocalBoxNotify")
}
}
pub trait Notify {
type Event;
fn poll_next(self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<Self::Event>;
}
impl<N> Notify for Box<N>
where
N: ?Sized + Notify + Unpin,
{
type Event = N::Event;
#[inline]
fn poll_next(self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<N::Event> {
Pin::new(self.get_mut().as_mut()).poll_next(t)
}
}
impl<N, P> Notify for Pin<P>
where
P: core::ops::DerefMut<Target = N> + Unpin,
N: Notify + ?Sized,
{
type Event = N::Event;
fn poll_next(self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<Self::Event> {
Pin::get_mut(self).as_mut().poll_next(t)
}
}
impl<N> Notify for &mut N
where
N: Notify + Unpin + ?Sized,
{
type Event = N::Event;
#[inline]
fn poll_next(mut self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<N::Event> {
Pin::new(&mut **self).poll_next(t)
}
}
impl<N> Notify for [N]
where
N: Notify + Unpin,
{
type Event = (usize, N::Event);
#[inline]
fn poll_next(self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<Self::Event> {
for (i, this) in self.get_mut().iter_mut().enumerate() {
if let Poll::Ready(value) = Pin::new(this).poll_next(t) {
return Poll::Ready((i, value));
}
}
Poll::Pending
}
}
pub trait NotifyExt: Notify + Sized + Unpin {
#[inline(always)]
fn next(&mut self) -> Next<'_, Self> {
Next(self)
}
#[inline(always)]
fn map<F>(self, f: F) -> Map<Self, F> {
let noti = self;
Map { noti, f }
}
}
impl<N: Notify + Sized + Unpin> NotifyExt for N {}
#[derive(Debug)]
pub struct Next<'a, N>(&'a mut N)
where
N: Notify + Unpin;
impl<N> Future for Next<'_, N>
where
N: Notify + Unpin,
{
type Output = N::Event;
#[inline]
fn poll(self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<Self::Output> {
Pin::new(&mut self.get_mut().0).poll_next(t)
}
}
pub trait Fuse: Sized {
fn fuse(self) -> Option<Self>;
}
impl<F> Fuse for F
where
F: Future,
{
fn fuse(self) -> Option<Self> {
self.into()
}
}
impl<F: Future> Notify for Option<F> {
type Event = F::Output;
#[inline]
fn poll_next(self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<F::Output> {
let mut s = self;
let out = s.as_mut().as_pin_mut().map(|f| f.poll(t));
if matches!(out, Some(Poll::Ready(_))) {
s.set(None);
}
out.unwrap_or(Poll::Pending)
}
}
#[derive(Debug)]
pub struct Map<N, F> {
noti: N,
f: F,
}
impl<N, F, E> Notify for Map<N, F>
where
N: Notify + Unpin,
F: FnMut(N::Event) -> E + Unpin,
{
type Event = E;
#[inline]
fn poll_next(mut self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<E> {
Pin::new(&mut self.noti).poll_next(t).map(&mut self.f)
}
}
#[derive(Debug)]
pub struct FutureFn<T, F>(T, F);
impl<T, F> Notify for FutureFn<T, F>
where
T: Future + Unpin,
F: FnMut() -> T + Unpin,
{
type Event = T::Output;
#[inline]
fn poll_next(self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<T::Output> {
let this = self.get_mut();
let poll = Pin::new(&mut this.0).poll(t);
if poll.is_ready() {
Pin::new(&mut this.0).set(this.1());
}
poll
}
}
#[derive(Debug)]
pub struct PollFn<F>(F);
impl<T, F> Notify for PollFn<F>
where
F: FnMut(&mut Task<'_>) -> Poll<T> + Unpin,
{
type Event = T;
#[inline]
fn poll_next(self: Pin<&mut Self>, t: &mut Task<'_>) -> Poll<T> {
self.get_mut().0(t)
}
}
#[derive(Debug)]
pub struct Pending<T>(core::marker::PhantomData<fn() -> T>);
impl<T> Notify for Pending<T> {
type Event = T;
fn poll_next(self: Pin<&mut Self>, _task: &mut Task<'_>) -> Poll<T> {
Poll::Pending
}
}
#[derive(Debug)]
pub struct Ready<T: Unpin>(Option<T>);
impl<T: Unpin> Notify for Ready<T> {
type Event = T;
fn poll_next(self: Pin<&mut Self>, _task: &mut Task<'_>) -> Poll<T> {
let Some(event) = self.get_mut().0.take() else {
return Poll::Pending;
};
Poll::Ready(event)
}
}
pub struct Select<'a, E, const N: usize>(
[&'a mut (dyn Notify<Event = E> + Unpin); N],
usize,
);
impl<E, const N: usize> fmt::Debug for Select<'_, E, N> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("Select")
}
}
impl<E: Unpin, const N: usize> Notify for Select<'_, E, N> {
type Event = E;
fn poll_next(self: Pin<&mut Self>, task: &mut Task<'_>) -> Poll<E> {
let s = self.get_mut();
let start = s.1;
if N == 0 {
return Poll::Pending;
}
for i in (start..N).chain(0..start) {
if let Poll::Ready(event) = Pin::new(&mut s.0[i]).poll_next(task) {
return Poll::Ready(event);
}
}
Poll::Pending
}
}
pub fn future_fn<T, F>(mut f: F) -> FutureFn<T, F>
where
T: Future + Unpin,
F: FnMut() -> T + Unpin,
{
FutureFn(f(), f)
}
pub fn poll_fn<T, F>(f: F) -> PollFn<F>
where
F: FnMut(&mut Task<'_>) -> Poll<T> + Unpin,
{
PollFn(f)
}
pub fn pending<T>() -> Pending<T> {
Pending(core::marker::PhantomData)
}
pub fn ready<T: Unpin>(t: T) -> Ready<T> {
Ready(t.into())
}
pub fn select<E, const N: usize>(
notifys: [&mut (dyn Notify<Event = E> + Unpin); N],
) -> Select<'_, E, N> {
Select(notifys, 0)
}