embedded-draw-target 0.1.1

//! ISR-safe DMA completion signaling for async and sync present paths.

use core::cell::Cell;
use core::future::Future;
use core::pin::Pin;
use core::sync::atomic::{AtomicBool, Ordering};
use core::task::{Context, Poll, Waker};

use critical_section::Mutex;
use embedded_graphics_core::pixelcolor::Rgb565;
use embedded_graphics_framebuf::{FrameBuf, backends::DMACapableFrameBufferBackend};

/// Trait representing a DMA transfer token.
pub trait DmaTransfer {
    /// The buffer type returned when the transfer finishes.
    type Buffer;

    /// Returns `true` if the DMA transfer has completed.
    fn is_done(&self) -> bool;

    /// Block until the transfer completes, recovering ownership of the buffer.
    fn wait(self) -> Self::Buffer;
}

/// Trait representing an async DMA transfer token.
pub trait AsyncDmaTransfer: DmaTransfer {
    /// The future type returned by [`wait_async`](Self::wait_async).
    type WaitFuture: Future<Output = Self::Buffer>;

    /// Return a future that resolves when the DMA transfer completes.
    fn wait_async(self) -> Self::WaitFuture;
}

/// One-shot completion flag with optional async waker notification.
pub struct CompletionSlot {
    signaled: AtomicBool,
    waker: Mutex<Cell<Option<Waker>>>,
}

impl CompletionSlot {
    /// Create a cleared completion slot.
    pub const fn new() -> Self {
        Self {
            signaled: AtomicBool::new(false),
            waker: Mutex::new(Cell::new(None)),
        }
    }

    /// Clear the slot before kicking off a new DMA transfer.
    pub fn reset(&self) {
        self.signaled.store(false, Ordering::Release);
        critical_section::with(|cs| {
            self.waker.borrow(cs).set(None);
        });
    }

    /// Mark complete and wake any registered async task.
    pub fn signal(&self) {
        self.signaled.store(true, Ordering::Release);
        critical_section::with(|cs| {
            if let Some(waker) = self.waker.borrow(cs).take() {
                waker.wake();
            }
        });
    }

    /// Returns `true` after [`signal`](Self::signal) until the next [`reset`](Self::reset).
    pub fn is_signaled(&self) -> bool {
        self.signaled.load(Ordering::Acquire)
    }

    /// Poll for completion, registering `cx`'s waker when still pending.
    pub fn poll_wait(&self, cx: &mut Context<'_>) -> Poll<()> {
        if self.is_signaled() {
            return Poll::Ready(());
        }
        critical_section::with(|cs| {
            self.waker.borrow(cs).set(Some(cx.waker().clone()));
        });
        if self.is_signaled() {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }
}

impl Default for CompletionSlot {
    fn default() -> Self {
        Self::new()
    }
}

/// Reference [`AsyncDmaTransfer`] token backed by a [`CompletionSlot`].
pub struct WaitTransfer<FB>
where
    FB: DMACapableFrameBufferBackend<Color = Rgb565>,
{
    framebuffer: Option<FrameBuf<Rgb565, FB>>,
    completion: &'static CompletionSlot,
}

impl<FB> WaitTransfer<FB>
where
    FB: DMACapableFrameBufferBackend<Color = Rgb565>,
{
    /// Build a transfer token; `completion` must outlive all in-flight DMA ops.
    pub fn new(framebuffer: FrameBuf<Rgb565, FB>, completion: &'static CompletionSlot) -> Self {
        completion.reset();
        Self {
            framebuffer: Some(framebuffer),
            completion,
        }
    }

    /// The completion slot wired to this transfer.
    pub const fn completion(&self) -> &'static CompletionSlot {
        self.completion
    }
}

impl<FB> DmaTransfer for WaitTransfer<FB>
where
    FB: DMACapableFrameBufferBackend<Color = Rgb565>,
{
    type Buffer = FrameBuf<Rgb565, FB>;

    fn is_done(&self) -> bool {
        self.completion.is_signaled()
    }

    fn wait(self) -> Self::Buffer {
        while !self.completion.is_signaled() {
            core::hint::spin_loop();
        }
        self.framebuffer
            .expect("WaitTransfer polled after completion")
    }
}

/// Future returned by [`WaitTransfer::wait_async`].
pub struct WaitTransferFuture<FB>
where
    FB: DMACapableFrameBufferBackend<Color = Rgb565>,
{
    inner: Option<WaitTransfer<FB>>,
}

impl<FB> AsyncDmaTransfer for WaitTransfer<FB>
where
    FB: DMACapableFrameBufferBackend<Color = Rgb565>,
{
    type WaitFuture = WaitTransferFuture<FB>;

    fn wait_async(self) -> Self::WaitFuture {
        WaitTransferFuture { inner: Some(self) }
    }
}

impl<FB> Future for WaitTransferFuture<FB>
where
    FB: DMACapableFrameBufferBackend<Color = Rgb565>,
{
    type Output = FrameBuf<Rgb565, FB>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = unsafe { self.get_unchecked_mut() };
        let inner = this
            .inner
            .as_mut()
            .expect("WaitTransferFuture polled after completion");
        match inner.completion.poll_wait(cx) {
            Poll::Ready(()) => Poll::Ready(
                this.inner
                    .take()
                    .expect("WaitTransferFuture polled after completion")
                    .framebuffer
                    .expect("WaitTransferFuture polled after completion"),
            ),
            Poll::Pending => Poll::Pending,
        }
    }
}