no-alloc-channel 0.1.0

#![cfg_attr(not(test), no_std)]
//! A wrapper of [`Queue`] which allows for asynchronous receiving.

use core::{
    future::Future,
    marker::PhantomPinned,
    pin::Pin,
    ptr::null_mut,
    sync::atomic::{AtomicPtr, Ordering},
    task::{Context, Poll, Waker},
};

use heapless::spsc::{Consumer, Producer, Queue};

/// A thread-safe channel supporting `async` receiving, with a buffer of up to
/// `N - 1` elements of type `T`.
pub struct Channel<T: Sync + 'static, const N: usize> {
    queue: Queue<T, N>,
    waker: AtomicPtr<Waker>,
}

impl<T: Send + Sync + 'static, const N: usize> Channel<T, N> {
    /// Construct a new, empty.
    pub const fn new() -> Self {
        Self {
            queue: Queue::new(),
            waker: AtomicPtr::new(null_mut()),
        }
    }

    /// Split the channel into a [`SendHalf`] and a [`ReceiveHalf`].
    pub fn split(&mut self) -> (SendHalf<'_, T, N>, ReceiveHalf<'_, T, N>) {
        let (producer, consumer) = self.queue.split();
        let waker = &self.waker;
        (
            SendHalf { producer, waker },
            ReceiveHalf { consumer, waker },
        )
    }
}

/// The sending half of a [`Channel`].
pub struct SendHalf<'a, T: 'static, const N: usize> {
    producer: Producer<'a, T, N>,
    waker: &'a AtomicPtr<Waker>,
}

impl<'a, T: 'static, const N: usize> SendHalf<'a, T, N> {
    /// Try to send a value on the channel, returning [`Err`] if the channel is
    /// full.
    pub fn try_send(&mut self, packet: T) -> Result<(), SendError> {
        self.producer.enqueue(packet).map_err(|_| SendError::Full)?;
        if let Some(waker) = unsafe { self.waker.load(Ordering::Acquire).as_ref() } {
            waker.wake_by_ref();
        }
        Ok(())
    }
}

/// The receiving half of a [`Channel`].
pub struct ReceiveHalf<'a, T: Send + 'static, const N: usize> {
    consumer: Consumer<'a, T, N>,
    waker: &'a AtomicPtr<Waker>,
}

impl<'a, T: Send + 'static, const N: usize> ReceiveHalf<'a, T, N> {
    /// Obtain a future which resolves with a value received from the channel.
    ///
    /// The receiving operation is atomic, and will only actually remove the
    /// element from the channel in the process of resolving to a
    /// [`Poll::Ready`] value.
    pub fn receive(&mut self) -> impl Future<Output = T> + Send + '_ {
        struct ReceiveFuture<'a, 'b, T: Send + 'static, const N: usize> {
            receive: &'a mut ReceiveHalf<'b, T, N>,
            waker: Option<Waker>,
            _phantom: PhantomPinned,
        }

        impl<'a, 'b, T: Send + 'static, const N: usize> Drop for ReceiveFuture<'a, 'b, T, N> {
            fn drop(&mut self) {
                // The waker pointer is known to currently be either null or point to the waker
                // inside self.waker (if that is Some(...)). Therefore, no matter what, we want
                // it to be null after this.
                self.receive.waker.store(null_mut(), Ordering::Release);
            }
        }

        impl<'a, 'b, T: Send + 'static, const N: usize> Future for ReceiveFuture<'a, 'b, T, N> {
            type Output = T;

            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
                let self_mut = unsafe { self.get_unchecked_mut() };
                if let Some(value) = self_mut.receive.try_receive() {
                    Poll::Ready(value)
                } else {
                    if self_mut.waker.is_none() {
                        let waker = self_mut.waker.get_or_insert(cx.waker().clone());
                        self_mut.receive.waker.store(waker, Ordering::Release);
                    }
                    Poll::Pending
                }
            }
        }

        ReceiveFuture {
            receive: self,
            waker: None,
            _phantom: PhantomPinned,
        }
    }

    /// Try to receive a value, returning [`None`] if the channel is empty.
    pub fn try_receive(&mut self) -> Option<T> {
        self.consumer.dequeue()
    }
}

/// An error that can arise from sending on a channel.
#[derive(Debug)]
pub enum SendError {
    /// The channel is full.
    Full,
}

#[cfg(test)]
#[tokio::test]
async fn simple() {
    let mut ch: Channel<u32, 4> = Channel::new();
    let (mut tx, mut rx) = ch.split();

    assert!(matches!(tx.try_send(0), Ok(())));
    assert!(matches!(tx.try_send(1), Ok(())));
    assert!(matches!(tx.try_send(2), Ok(())));
    assert!(matches!(tx.try_send(3), Err(SendError::Full)));

    assert_eq!(rx.receive().await, 0);
    assert_eq!(rx.receive().await, 1);
    assert_eq!(rx.receive().await, 2);
    assert!(matches!(rx.try_receive(), None));
}