hirun 0.1.7

use super::{sched_from_ctx, task_from_ctx, TaskRef};
use crate::event::{Scheduler, Timer};
use core::future::Future;
use core::mem::MaybeUninit;
use core::pin::Pin;
use core::task::{Context, Poll};
use core::time::Duration;
use hioff::container_of_mut;

/// 通过宏`#[future]`使用。目的时消除Rust自动判定的Send存在的约束.
/// 一个Future内部只是多个异步函数串行，也无法跨子异步函数的await使用不支持Send的数据类型，比如不能使用Rc,只能使用Arc.
/// 另外因为异步函数的特殊性，无法类似thread::scope来消除生命周期相关的使用约束。这些实际功能正确无风险的代码可能导致Future无法满足Send和`'static`的约束, 此时可以用`#[future]`来消除这类约束.
pub struct SendFuture<T: Future> {
    future: T,
}

unsafe impl<T: Future> Send for SendFuture<T> {}

impl<T: Future> SendFuture<T> {
    /// # Safety
    /// 使用人员确保`#[future]`的正确使用，future内部不支持Send的数据类型不存在并发访问场景.
    pub unsafe fn new(future: T) -> Self {
        Self { future }
    }
}

impl<T: Future> Future for SendFuture<T> {
    type Output = T::Output;
    #[inline(always)]
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        Future::poll(
            unsafe { Pin::new_unchecked(&mut Pin::into_inner_unchecked(self).future) },
            ctx,
        )
    }
}

pub(crate) struct FnOnceFuture<F, R>(MaybeUninit<F>)
where
    F: FnOnce() -> R;

impl<F, R> FnOnceFuture<F, R>
where
    F: FnOnce() -> R,
{
    pub(crate) fn new(f: F) -> Self {
        Self(MaybeUninit::new(f))
    }
}

impl<F, R> Future for FnOnceFuture<F, R>
where
    F: FnOnce() -> R,
{
    type Output = R;
    fn poll(self: Pin<&mut Self>, _ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = unsafe { &mut Pin::into_inner_unchecked(self) };
        Poll::Ready(unsafe { this.0.assume_init_read() }())
    }
}

pub(crate) struct Sleep {
    timer: Timer,
    task: Option<TaskRef>,
    timeout: Duration,
}

impl Unpin for Sleep {}
unsafe impl Send for Sleep {}

impl Sleep {
    pub(crate) fn new(mut timeout: Duration) -> Self {
        if timeout > Duration::MAX - Duration::new(2, 0) {
            timeout = Duration::MAX - Duration::new(2, 0);
        }
        Self {
            timer: Timer::new(Self::timer_handle),
            task: None,
            timeout,
        }
    }

    fn timer_handle(timer: &Timer, sched: &mut Scheduler) {
        let this = unsafe { container_of_mut!(timer, Self, timer) };
        let task = this.task.take().unwrap();
        task.sched_waked(sched);
        this.timeout = Duration::MAX;
    }
}

impl Future for Sleep {
    type Output = ();
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.timeout == Duration::MAX {
            return Poll::Ready(());
        }
        if self.timeout < Duration::MAX - Duration::new(1, 0) {
            let mut task = unsafe { task_from_ctx(ctx).as_ref() }.task_ref();
            //如果过程中被abort，也只能在timer_handle之后才真正结束.
            task.status.inc_wake();
            task.status.set_yield();
            self.task = Some(task);
            let sched = unsafe { sched_from_ctx(ctx).as_mut() };
            unsafe { sched.set_timer(&self.timer, self.timeout.as_micros() as u32) };
            self.timeout = Duration::MAX - Duration::new(1, 0);
        }
        Poll::Pending
    }
}

pub(crate) struct Yield(bool);

impl Yield {
    pub fn new() -> Self {
        Self(false)
    }
}

impl Future for Yield {
    type Output = ();
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.0 {
            return Poll::Ready(());
        }
        let task = unsafe { task_from_ctx(ctx).as_mut() };
        task.status.inc_wake();
        task.status.set_yield();
        let sched = unsafe { sched_from_ctx(ctx).as_mut() };
        task.task_ref().sched_waked(sched);
        self.0 = true;
        Poll::Pending
    }
}

pub(crate) struct Delay<T: Future> {
    future: T,
    sleep: Sleep,
}

impl<T: Future> Unpin for Delay<T> {}
unsafe impl<T: Future + Send> Send for Delay<T> {}

impl<T: Future> Delay<T> {
    pub fn new(delay: Duration, future: T) -> Self {
        Self {
            sleep: Sleep::new(delay),
            future,
        }
    }
}

impl<T: Future> Future for Delay<T> {
    type Output = T::Output;
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let sleep = unsafe { Pin::new_unchecked(&mut self.sleep) };
        if Future::poll(sleep, ctx) == Poll::Ready(()) {
            let pinned = unsafe { Pin::new_unchecked(&mut self.future) };
            return Future::poll(pinned, ctx);
        }
        Poll::Pending
    }
}

pub(crate) struct Deadline<T: Future> {
    future: T,
    timeout: Duration,
}

impl<T: Future> Unpin for Deadline<T> {}
unsafe impl<T: Future + Send> Send for Deadline<T> {}

impl<T: Future> Deadline<T> {
    pub fn new(timeout: Duration, future: T) -> Self {
        Self { future, timeout }
    }
}

impl<T: Future> Future for Deadline<T> {
    type Output = T::Output;
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let tm = crate::time::now();
        if self.timeout >= tm {
            let pinned = unsafe { Pin::new_unchecked(&mut self.future) };
            return Future::poll(pinned, ctx);
        }

        let task = unsafe { task_from_ctx(ctx).as_mut() };
        task.poll_deadline();
        task.status.inc_wake();
        Poll::Pending
    }
}