async-utils 0.1.1

use std::{
    cell::{Cell, UnsafeCell},
    fmt,
    future::Future,
    marker::PhantomPinned,
    mem,
    ops::{Deref, DerefMut},
    panic::{RefUnwindSafe, UnwindSafe},
    pin::Pin,
    ptr::{self, NonNull},
    sync::atomic::{AtomicUsize, Ordering},
    task::{Context, Poll, Waker},
};

use crate::sync::{
    atomic::full_fence,
    mutex::{Mutex, MutexFamily, MutexGuardType, StdMutex},
};

pub struct Event<M: MutexFamily = StdMutex> {
    /// The number of notified entries, or `usize::MAX` if all of them have been notified.
    ///
    /// If there are no entries, this value is set to `usize::MAX`.
    notified: AtomicUsize,

    /// A linked list holding registered listeners.
    list: M::Mutex<List>,
}

unsafe impl<M: MutexFamily> Send for Event<M> {}
unsafe impl<M: MutexFamily> Sync for Event<M> {}

impl<M: MutexFamily> UnwindSafe for Event<M> {}
impl<M: MutexFamily> RefUnwindSafe for Event<M> {}

impl<M: MutexFamily> Event<M> {
    #[inline]
    pub fn new() -> Self {
        Self {
            notified: AtomicUsize::new(usize::MAX),
            list: M::new(List::new()),
        }
    }

    #[inline]
    fn lock(&self) -> ListGuard<'_, M> {
        ListGuard {
            event: self,
            guard: self.list.lock(),
        }
    }

    #[cold]
    pub fn listener(&self) -> EventListener<'_, M> {
        EventListener {
            event: self,
            state: ListenerState::Init,
            entry: UnsafeCell::new(Entry::new()),
        }
    }

    #[inline]
    pub fn notify(&self, n: usize) {
        // Make sure the notification comes after whatever triggered it.
        full_fence();
        self.notify_relaxed(n)
    }

    #[inline]
    pub fn notify_relaxed(&self, n: usize) {
        // Notify if there is at least one unnotified listener and the number of notified
        // listeners is less than `n`.
        if self.notified.load(Ordering::Acquire) < n {
            self.lock().notify(n);
        }
    }

    #[inline]
    pub fn notify_additional(&self, n: usize) {
        // Make sure the notification comes after whatever triggered it.
        full_fence();
        self.notify_additional_relaxed(n);
    }

    #[inline]
    pub fn notify_additional_relaxed(&self, n: usize) {
        // Notify if there is at least one unnotified listener.
        if self.notified.load(Ordering::Acquire) < usize::MAX {
            self.lock().notify_additional(n);
        }
    }
}

impl<M: MutexFamily> fmt::Debug for Event<M> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.pad("Event { .. }")
    }
}

impl<M: MutexFamily> Default for Event<M> {
    fn default() -> Self {
        Event::new()
    }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
enum ListenerState {
    Init,
    Listening,
    Done,
}

pub struct EventListener<'a, M: MutexFamily> {
    event: &'a Event<M>,
    state: ListenerState,
    entry: UnsafeCell<Entry>,
}

unsafe impl<M: MutexFamily> Send for EventListener<'_, M> {}
unsafe impl<M: MutexFamily> Sync for EventListener<'_, M> {}

impl<M: MutexFamily> UnwindSafe for EventListener<'_, M> {}
impl<M: MutexFamily> RefUnwindSafe for EventListener<'_, M> {}

impl<M: MutexFamily> EventListener<'_, M> {
    #[inline]
    fn project(self: Pin<&mut Self>) -> (&Event<M>, &mut ListenerState, Pin<&UnsafeCell<Entry>>) {
        fn is_unpin<T: Unpin>() {}

        unsafe {
            is_unpin::<&Event<M>>();
            is_unpin::<ListenerState>();

            let this = self.get_unchecked_mut();
            (this.event, &mut this.state, Pin::new_unchecked(&this.entry))
        }
    }

    pub fn listen(self: Pin<&mut Self>) {
        let (event, state, entry) = self.project();

        if *state == ListenerState::Init {
            event.lock().insert(entry, None);
            *state = ListenerState::Listening;
        }

        full_fence();
    }

    #[inline]
    pub fn listens_to(&self, event: &Event<M>) -> bool {
        ptr::eq(self.event, event)
    }

    #[inline]
    pub fn same_event(&self, other: &EventListener<'_, M>) -> bool {
        ptr::eq(self.event, other.event)
    }
}

impl<M: MutexFamily> Future for EventListener<'_, M> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let (inner, listener_state, entry) = self.project();

        if *listener_state == ListenerState::Done {
            return Poll::Ready(());
        }

        let mut list = inner.lock();

        if *listener_state == ListenerState::Init {
            list.insert(entry, Some(cx.waker().clone()));
            *listener_state = ListenerState::Listening;

            full_fence();

            return Poll::Pending;
        }

        let entry = unsafe { &mut *entry.get_ref().get() };
        let state = &mut entry.state;

        // Do a dummy replace operation in order to take out the state.
        match state.replace(State::Notified(false)) {
            State::Notified(_) => {
                // If this listener has been notified, remove it from the list and return.
                let entry = unsafe { NonNull::new_unchecked(entry as *mut _) };
                list.remove(entry);
                *listener_state = ListenerState::Done;

                drop(list);

                return Poll::Ready(());
            }
            State::Created => {
                // If the listener was just created, put it in the `Polling` state.
                state.set(State::Polling(cx.waker().clone()));
            }
            State::Polling(w) => {
                // If the listener was in the `Polling` state, update the waker.
                if w.will_wake(cx.waker()) {
                    state.set(State::Polling(w));
                } else {
                    state.set(State::Polling(cx.waker().clone()));
                }
            }
        }

        Poll::Pending
    }
}

impl<M: MutexFamily> Drop for EventListener<'_, M> {
    fn drop(&mut self) {
        // If this listener has never picked up a notification...
        if self.state == ListenerState::Listening {
            let mut list = self.event.lock();
            let entry = unsafe { NonNull::new_unchecked(self.entry.get()) };

            // But if a notification was delivered to it...
            if let State::Notified(additional) = list.remove(entry) {
                // Then pass it on to another active listener.
                if additional {
                    list.notify_additional(1);
                } else {
                    list.notify(1);
                }
            }

            self.state = ListenerState::Done;
        }
    }
}

/// A guard holding the linked list locked.
struct ListGuard<'a, M: MutexFamily> {
    /// A reference to [`Event`]'s state.
    event: &'a Event<M>,

    /// The actual guard that acquired the linked list.
    guard: MutexGuardType<'a, M, List>,
}

impl<M: MutexFamily> Drop for ListGuard<'_, M> {
    #[inline]
    fn drop(&mut self) {
        let list = &mut **self;

        // Update the atomic `notified` counter.
        let notified = if list.notified < list.len {
            list.notified
        } else {
            usize::MAX
        };

        self.event.notified.store(notified, Ordering::Release);
    }
}

impl<M: MutexFamily> Deref for ListGuard<'_, M> {
    type Target = List;

    #[inline]
    fn deref(&self) -> &List {
        &self.guard
    }
}

impl<M: MutexFamily> DerefMut for ListGuard<'_, M> {
    #[inline]
    fn deref_mut(&mut self) -> &mut List {
        &mut self.guard
    }
}

/// The state of a listener.
enum State {
    /// It has just been created.
    Created,

    /// It has received a notification.
    ///
    /// The `bool` is `true` if this was an "additional" notification.
    Notified(bool),

    /// An async task is polling it.
    Polling(Waker),
}

impl State {
    /// Returns `true` if this is the `Notified` state.
    #[inline]
    fn is_notified(&self) -> bool {
        match self {
            State::Notified(_) => true,
            State::Created | State::Polling(_) => false,
        }
    }
}

/// An entry representing a registered listener.
struct Entry {
    /// THe state of this listener.
    state: Cell<State>,

    /// Previous entry in the linked list.
    prev: Cell<Option<NonNull<Entry>>>,

    /// Next entry in the linked list.
    next: Cell<Option<NonNull<Entry>>>,

    /// Entries are pinned in memory.
    _pinned: PhantomPinned,
}

impl Entry {
    #[inline]
    const fn new() -> Self {
        Self {
            state: Cell::new(State::Created),
            prev: Cell::new(None),
            next: Cell::new(None),
            _pinned: PhantomPinned,
        }
    }
}

/// A linked list of entries.
struct List {
    /// First entry in the list.
    head: Option<NonNull<Entry>>,

    /// Last entry in the list.
    tail: Option<NonNull<Entry>>,

    /// The first unnotified entry in the list.
    start: Option<NonNull<Entry>>,

    /// Total number of entries in the list.
    len: usize,

    /// The number of notified entries in the list.
    notified: usize,
}

unsafe impl Send for List {}

impl List {
    #[inline]
    const fn new() -> Self {
        Self {
            head: None,
            tail: None,
            start: None,
            len: 0,
            notified: 0,
        }
    }

    /// Inserts a pinned entry into the list.
    fn insert(&mut self, entry: Pin<&UnsafeCell<Entry>>, waker: Option<Waker>) {
        unsafe {
            let state = waker.map(State::Polling).unwrap_or(State::Created);
            let entry = &mut *entry.get_ref().get();

            *entry = Entry {
                state: Cell::new(state),
                prev: Cell::new(self.tail),
                next: Cell::new(None),
                _pinned: PhantomPinned,
            };

            let entry = NonNull::new_unchecked(entry as *mut _);

            // Replace the tail with the new entry.
            match mem::replace(&mut self.tail, Some(entry)) {
                None => self.head = Some(entry),
                Some(t) => t.as_ref().next.set(Some(entry)),
            }

            // If there were no unnotified entries, this one is the first now.
            if self.start.is_none() {
                self.start = self.tail;
            }

            // Bump the entry count.
            self.len += 1;
        }
    }

    /// Removes an entry from the list and returns its state.
    fn remove(&mut self, entry: NonNull<Entry>) -> State {
        unsafe {
            let prev = entry.as_ref().prev.get();
            let next = entry.as_ref().next.get();

            // Unlink from the previous entry.
            match prev {
                None => self.head = next,
                Some(p) => p.as_ref().next.set(next),
            }

            // Unlink from the next entry.
            match next {
                None => self.tail = prev,
                Some(n) => n.as_ref().prev.set(prev),
            }

            // If this was the first unnotified entry, move the pointer to the next one.
            if self.start == Some(entry) {
                self.start = next;
            }

            // Extract the state.
            let state = entry.as_ref().state.replace(State::Created);

            // Update the counters.
            if state.is_notified() {
                self.notified -= 1;
            }
            self.len -= 1;

            state
        }
    }

    /// Notifies a number of entries.
    #[cold]
    fn notify(&mut self, mut n: usize) {
        if n <= self.notified {
            return;
        }
        n -= self.notified;

        while n > 0 {
            n -= 1;

            // Notify the first unnotified entry.
            match self.start {
                None => break,
                Some(e) => {
                    // Get the entry and move the pointer forward.
                    let e = unsafe { e.as_ref() };
                    self.start = e.next.get();

                    // Set the state of this entry to `Notified` and notify.
                    match e.state.replace(State::Notified(false)) {
                        State::Notified(_) => {}
                        State::Created => {}
                        State::Polling(w) => w.wake(),
                    }

                    // Update the counter.
                    self.notified += 1;
                }
            }
        }
    }

    /// Notifies a number of additional entries.
    #[cold]
    fn notify_additional(&mut self, mut n: usize) {
        while n > 0 {
            n -= 1;

            // Notify the first unnotified entry.
            match self.start {
                None => break,
                Some(e) => {
                    // Get the entry and move the pointer forward.
                    let e = unsafe { e.as_ref() };
                    self.start = e.next.get();

                    // Set the state of this entry to `Notified` and notify.
                    match e.state.replace(State::Notified(true)) {
                        State::Notified(_) => {}
                        State::Created => {}
                        State::Polling(w) => w.wake(),
                    }

                    // Update the counter.
                    self.notified += 1;
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::sync::notify::Event;

    fn is_send<T: Send>() {}
    fn is_sync<T: Sync>() {}

    #[test]
    fn event_send_sync() {
        is_send::<Event>();
        is_sync::<Event>();
    }
}