hirun 0.1.16

use super::{RawTask, RawTaskContext, Worker};
use crate::{Error, Result};
use core::cell::Cell;
use core::future::Future;
use core::marker::PhantomData;
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::ptr;
use core::sync::atomic::{
    AtomicPtr,
    Ordering::{Acquire, Relaxed, Release},
};
use core::task::{Context, Poll};

pub struct Mutex<T> {
    waiting: AtomicPtr<RawTask>,
    task: Cell<*mut RawTask>,
    val: T,
}

unsafe impl<T: Send> Send for Mutex<T> {}
unsafe impl<T> Sync for Mutex<T> {}

impl<T> Mutex<T> {
    pub fn new(val: T) -> Self {
        Self {
            waiting: AtomicPtr::<RawTask>::new(ptr::null_mut::<RawTask>()),
            task: Cell::new(ptr::null_mut::<RawTask>()),
            val,
        }
    }
}

impl<T> Mutex<T> {
    pub fn lock(&self) -> impl Future<Output = MutexGuard<'_, T>> {
        MutexFuture(self)
    }

    pub fn lock_blocking(&self) -> Result<MutexGuard<'_, T>> {
        crate::runtime::block_on(MutexFuture(self))
    }

    pub fn try_lock(&self) -> Result<MutexGuard<'_, T>> {
        if let Some(worker) = Worker::current() {
            let task = unsafe { worker.current_task().as_mut() };
            let next = ptr::null_mut::<RawTask>();
            unsafe { task.node.set_next(next) };
            if self
                .waiting
                .compare_exchange_weak(next, task, Release, Acquire)
                .is_ok()
            {
                self.task.set(task);
                return Ok(MutexGuard::new(self));
            }
        }
        Err(Error::default())
    }
}

pub struct MutexFuture<'a, T>(&'a Mutex<T>);

unsafe impl<T> Send for MutexFuture<'_, T> {}

impl<'a, T> Future for MutexFuture<'a, T> {
    type Output = MutexGuard<'a, T>;
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let task = ctx.task();
        if ptr::eq(self.0.task.get(), task) {
            return Poll::Ready(MutexGuard::new(self.get_mut().0));
        }

        let mut next = self.0.waiting.load(Acquire);
        loop {
            unsafe { task.node.set_next(next) };
            match self
                .0
                .waiting
                .compare_exchange_weak(next, task, Release, Acquire)
            {
                Ok(_) => break,
                Err(old) => next = old,
            }
        }
        if next.is_null() {
            self.0.task.set(task);
            Poll::Ready(MutexGuard::new(self.get_mut().0))
        } else {
            Poll::Pending
        }
    }
}

pub struct MutexGuard<'a, T> {
    mutex: *mut Mutex<T>,
    mark: PhantomData<&'a mut Mutex<T>>,
}

unsafe impl<T> Send for MutexGuard<'_, T> {}

impl<'a, T> MutexGuard<'a, T> {
    fn new(mutex: &'a Mutex<T>) -> Self {
        Self {
            mutex: mutex as *const Mutex<T> as *mut Mutex<T>,
            mark: PhantomData,
        }
    }
}

impl<T> Deref for MutexGuard<'_, T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        unsafe { &(*self.mutex).val }
    }
}

impl<T> DerefMut for MutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut (*self.mutex).val }
    }
}

impl<T> Drop for MutexGuard<'_, T> {
    fn drop(&mut self) {
        let mutex = unsafe { &*self.mutex };
        let mut next = mutex.waiting.load(Acquire);
        let task = mutex.task.replace(ptr::null_mut::<RawTask>());
        if ptr::eq(next, task) {
            match mutex
                .waiting
                .compare_exchange(task, ptr::null_mut::<RawTask>(), Relaxed, Acquire)
            {
                Ok(_) => return,
                Err(task) => next = task,
            }
        }
        while !ptr::eq(unsafe { (*next).node.next() }, task) {
            next = unsafe { (*next).node.next() };
        }
        unsafe { (*next).node.init_next() };

        mutex.task.set(next);
        unsafe { &*next }.wake();
    }
}

#[cfg(test)]
mod test {
    use crate::runtime::*;

    #[test]
    fn test_mutex() {
        let _ = Builder::new().build();
        let mutex = Mutex::new(1);
        let mutex = unsafe { &*(&mutex as *const Mutex<i32>) };

        async fn foo(mutex: &Mutex<i32>) -> i32 {
            for _ in 0..2 {
                let mut guard = mutex.lock().await;
                *guard += 1
            }
            let guard = mutex.lock().await;
            *guard
        }

        let h1 = spawn(foo(mutex));
        let h2 = spawn(foo(mutex));
        let h3 = spawn(foo(mutex));
        let _ = h1.join().unwrap() + h2.join().unwrap() + h3.join().unwrap();
        assert_eq!(*mutex.lock_blocking().unwrap(), 7);
    }
}