dope-runtime 0.2.2

use std::{
    future::Future,
    ptr::NonNull,
    task::{Context, Poll, Waker},
    time::Instant,
};

use dope_core::driver::{CompletionBackend, CurrentSlot, DriverLifecycle, PoolDriver};
use dope_core::profile::{Profile, Throughput};

use super::ctx::{RuntimeCtx, RuntimeHandles};
use super::external::{ExternalRouter, current_external};
use super::join_handle::JoinHandle;
use super::{LocalExecutor, run_ready};
use crate::runtime::Runtime;

impl<B, P> Runtime<LocalExecutor<B, P>>
where
    B: PoolDriver,
    P: Profile,
{
    pub fn spawn<F>(&self, fut: F) -> JoinHandle<F::Output>
    where
        F: Future + 'static,
        F::Output: 'static,
    {
        self.reactor.exec.spawn(fut)
    }

    pub fn ctx(&mut self) -> RuntimeCtx<B> {
        let driver_ptr = NonNull::from(&mut self.driver);
        let extra = RuntimeHandles {
            exec: NonNull::from(&*self.reactor.exec),
            timers: NonNull::from(&*self.reactor.timers),
        };
        let extra_ptr = Box::leak(Box::new(extra)) as *const RuntimeHandles as *const ();
        RuntimeCtx::from_slot(CurrentSlot::<B>::new(driver_ptr, extra_ptr))
    }

    pub fn block_on<F>(&mut self, fut: F) -> F::Output
    where
        F: Future,
    {
        self.run_until(fut)
    }

    pub fn run_until<F>(&mut self, fut: F) -> F::Output
    where
        F: Future,
    {
        let mut fut = std::pin::pin!(fut);
        let waker = Waker::noop();
        let mut cx = Context::from_waker(waker);

        let Self {
            reactor, driver, ..
        } = self;
        let extra = RuntimeHandles {
            exec: NonNull::from(&*reactor.exec),
            timers: NonNull::from(&*reactor.timers),
        };
        let raw = NonNull::from(&mut *driver);
        let slot = CurrentSlot::<B>::new(raw, &extra as *const _ as *const ());
        let _guard = SlotGuard::<B>::install(slot);

        loop {
            reactor.timers.wake_due();
            run_ready::<P>(&reactor.exec);

            if let Poll::Ready(v) = fut.as_mut().poll(&mut cx) {
                return v;
            }

            let aux = current_external::<B>();

            let progressed = matches!(
                B::drive_with(driver, &mut ExternalRouter::<B> { aux }),
                Ok(true)
            );
            if let Some(a) = aux {
                unsafe { (a.pump)(a.state, driver) };
            }

            if progressed {
                reactor.timers.wake_due();
                run_ready::<P>(&reactor.exec);
                if let Poll::Ready(v) = fut.as_mut().poll(&mut cx) {
                    return v;
                }
            }

            if !reactor.exec.is_empty() {
                continue;
            }
            for _ in 0..4 {
                if let Some(a) = aux {
                    unsafe { (a.pump)(a.state, driver) };
                }
                reactor.timers.wake_due();
                run_ready::<P>(&reactor.exec);
                if let Poll::Ready(v) = fut.as_mut().poll(&mut cx) {
                    return v;
                }
                if !reactor.exec.is_empty() {
                    break;
                }
                if !driver.has_pending() {
                    break;
                }
            }
            if !reactor.exec.is_empty() {
                continue;
            }
            let now = Instant::now();
            if let Some(deadline) = reactor.timers.next_deadline() {
                let _ = driver.park_for(deadline.saturating_duration_since(now));
            } else if let Some(timeout) = P::IDLE_PARK {
                let _ = driver.park_for(timeout);
            } else {
                let _ = driver.park();
            }
        }
    }

    pub fn shutdown(&mut self) {
        let Self {
            reactor, driver, ..
        } = self;
        let extra = RuntimeHandles {
            exec: NonNull::from(&*reactor.exec),
            timers: NonNull::from(&*reactor.timers),
        };
        let raw = NonNull::from(&mut *driver);
        let slot = CurrentSlot::<B>::new(raw, &extra as *const _ as *const ());
        let _guard = SlotGuard::<B>::install(slot);

        let mut last_activity = Instant::now();
        loop {
            reactor.timers.wake_due();
            let ran_ready = run_ready::<P>(&reactor.exec);
            if ran_ready > 0 {
                last_activity = Instant::now();
            }
            let progressed = matches!(driver.drive(), Ok(true));
            if (ran_ready > 0 || progressed) && driver.has_pending() {
                continue;
            }

            let next_deadline = reactor.timers.next_deadline();
            if reactor.exec.is_empty() && next_deadline.is_none() {
                break;
            }

            let now = Instant::now();
            let timeout =
                reactor
                    .idle_park_policy
                    .timeout_for(now, next_deadline, None, last_activity);
            let Some(timeout) = timeout else {
                break;
            };
            let _ = driver.park_for(timeout);
        }
    }
}

pub fn block_on<B, F>(fut: F) -> F::Output
where
    B: PoolDriver,
    F: Future,
{
    let mut rt = Runtime::<LocalExecutor<B, Throughput>>::new_local()
        .expect("dope_runtime::block_on: failed to construct runtime");
    let out = rt.run_until(fut);
    rt.shutdown();
    out
}

pub fn current<B: CompletionBackend>() -> RuntimeCtx<B> {
    let slot = B::current_slot()
        .expect("dope_runtime runtime invariant violated: current runtime ctx missing");
    RuntimeCtx::from_slot(slot)
}

pub fn spawn<B, F>(fut: F) -> JoinHandle<F::Output>
where
    B: PoolDriver,
    F: Future + 'static,
    F::Output: 'static,
{
    current::<B>().spawn(fut)
}

pub(super) struct SlotGuard<B: CompletionBackend>(std::marker::PhantomData<B>);

impl<B: CompletionBackend> SlotGuard<B> {
    pub(super) fn install(slot: CurrentSlot<B>) -> Self {
        let prev = B::current_slot();
        assert!(
            prev.is_none(),
            "dope_runtime runtime invariant violated: runtime slot already set"
        );
        B::set_slot(Some(slot));
        Self(std::marker::PhantomData)
    }
}

impl<B: CompletionBackend> Drop for SlotGuard<B> {
    fn drop(&mut self) {
        B::set_slot(None);
    }
}