hipthread 0.1.3

use super::Mutex;
use core::cell::{Cell, UnsafeCell};
use core::fmt;
use core::mem::{ManuallyDrop, MaybeUninit};
use core::ops::{Deref, DerefMut};

/// 类似pthread_once的机制.
pub struct Once {
    // OnceLock/LazyLock需要获取更多的状态信息，直接使用pthread_once难以满足
    done: Cell<bool>,
    lock: Mutex<()>,
}

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

impl Once {
    pub const fn new() -> Self {
        Self {
            done: Cell::new(false),
            lock: Mutex::new(()),
        }
    }
    /// 确保只会初始化一次.
    pub fn call_once<F>(&self, f: F)
    where
        F: FnOnce(),
    {
        self.init(f);
    }

    fn init<F>(&self, f: F) -> bool
    where
        F: FnOnce(),
    {
        if self.done.get() {
            return false;
        }
        let _guard = self.lock.lock();
        if !self.done.get() {
            f();
            self.done.set(true);
            return true;
        }
        false
    }

    fn is_called(&self) -> bool {
        self.done.get()
    }

    fn set_called(&mut self, called: bool) {
        self.done.set(called);
    }
}

/// 类似std::sync::OnceLock
pub struct OnceLock<T> {
    value: UnsafeCell<MaybeUninit<T>>,
    once: Once,
}

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

impl<T> OnceLock<T> {
    pub const fn new() -> Self {
        Self {
            value: UnsafeCell::new(MaybeUninit::uninit()),
            once: Once::new(),
        }
    }

    pub fn get(&self) -> Option<&T> {
        if self.once.is_called() {
            return Some(unsafe { (*self.value.get()).assume_init_ref() });
        }
        None
    }

    pub fn get_mut(&mut self) -> Option<&mut T> {
        if self.once.is_called() {
            return Some(unsafe { (*self.value.get()).assume_init_mut() });
        }
        None
    }

    pub fn get_or_init<F>(&self, f: F) -> &T
    where
        F: FnOnce() -> T,
    {
        self.once.call_once(|| unsafe {
            (*self.value.get()).write(f());
        });
        unsafe { (*self.value.get()).assume_init_ref() }
    }

    pub fn get_mut_or_init<F>(&mut self, f: F) -> &mut T
    where
        F: FnOnce() -> T,
    {
        self.once.call_once(|| unsafe {
            (*self.value.get()).write(f());
        });
        unsafe { self.value.get_mut().assume_init_mut() }
    }

    pub fn take(&mut self) -> Option<T> {
        if self.once.is_called() {
            self.once.set_called(false);
            return Some(unsafe { self.value.get_mut().assume_init_read() });
        }
        None
    }

    pub fn set(&self, value: T) -> core::result::Result<(), T> {
        let mut value = ManuallyDrop::new(value);
        if self.once.init(|| unsafe {
            (*self.value.get()).write(ManuallyDrop::take(&mut value));
        }) {
            return Ok(());
        }
        Err(unsafe { ManuallyDrop::take(&mut value) })
    }
}

impl<T> Default for OnceLock<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Drop for OnceLock<T> {
    fn drop(&mut self) {
        if self.once.is_called() {
            unsafe { self.value.get_mut().assume_init_drop() };
        }
    }
}

impl<T: Clone> Clone for OnceLock<T> {
    fn clone(&self) -> Self {
        let cloned = Self::new();
        if let Some(value) = self.get() {
            let _ = cloned.set(value.clone());
        }
        cloned
    }
}

impl<T: fmt::Debug> fmt::Debug for OnceLock<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Debug::fmt(&self.get(), f)
    }
}

impl<T: fmt::Display> fmt::Display for OnceLock<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let Some(value) = self.get() {
            fmt::Display::fmt(value, f)
        } else {
            fmt::Display::fmt("<uninit>", f)
        }
    }
}

union LazyData<T, F> {
    data: ManuallyDrop<T>,
    f: ManuallyDrop<F>,
}

/// 类似std::sync::LazyLock
pub struct LazyLock<T, F = fn() -> T> {
    data: UnsafeCell<LazyData<T, F>>,
    once: Once,
}

unsafe impl<T: Sync, F> Sync for LazyLock<T, F> {}
unsafe impl<T: Send, F: Send> Send for LazyLock<T, F> {}

impl<T, F: FnOnce() -> T> LazyLock<T, F> {
    pub const fn new(f: F) -> Self {
        Self {
            once: Once::new(),
            data: UnsafeCell::new(LazyData {
                f: ManuallyDrop::new(f),
            }),
        }
    }

    pub fn force_mut(this: &mut Self) -> &mut T {
        this.init();
        unsafe { &mut this.data.get_mut().data }
    }

    pub fn force(this: &Self) -> &T {
        this.init();
        unsafe { &(*this.data.get()).data }
    }

    fn init(&self) {
        self.once.call_once(|| {
            let data = unsafe { &mut *self.data.get() };
            let f = unsafe { ManuallyDrop::take(&mut data.f) };
            data.data = ManuallyDrop::new(f());
        });
    }
}

impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        Self::force(self)
    }
}

impl<T, F: FnOnce() -> T> DerefMut for LazyLock<T, F> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        Self::force_mut(self)
    }
}
impl<T, F> Drop for LazyLock<T, F> {
    fn drop(&mut self) {
        if self.once.is_called() {
            unsafe { ManuallyDrop::drop(&mut self.data.get_mut().data) }
        } else {
            unsafe { ManuallyDrop::drop(&mut self.data.get_mut().f) }
        }
    }
}

impl<T: Default> Default for LazyLock<T> {
    fn default() -> Self {
        Self::new(T::default)
    }
}

impl<T: fmt::Debug, F: FnOnce() -> T> fmt::Debug for LazyLock<T, F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.once.is_called() {
            fmt::Debug::fmt(Self::force(self), f)
        } else {
            fmt::Debug::fmt("<uninit>", f)
        }
    }
}

impl<T: fmt::Display, F: FnOnce() -> T> fmt::Display for LazyLock<T, F> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.once.is_called() {
            fmt::Display::fmt(Self::force(self), f)
        } else {
            fmt::Display::fmt("<uninit>", f)
        }
    }
}