curiOS/kernel/src/sync/mod.rs

mod spin_mutex {
    use core::{
        cell::UnsafeCell,
        sync::atomic::{AtomicBool, Ordering},
    };

    pub struct SpinMutex<T: ?Sized> {
        lock: AtomicBool,
        data: UnsafeCell<T>,
    }

    pub struct SpinMutexGuard<'a, T: ?Sized + 'a> {
        mutex: &'a SpinMutex<T>,
    }

    impl<T: ?Sized> core::ops::Deref for SpinMutexGuard<'_, T> {
        type Target = T;

        fn deref(&self) -> &Self::Target {
            unsafe { &*self.mutex.data.get() }
        }
    }

    impl<T: ?Sized> core::ops::DerefMut for SpinMutexGuard<'_, T> {
        fn deref_mut(&mut self) -> &mut Self::Target {
            unsafe { &mut *self.mutex.data.get() }
        }
    }

    impl<T: ?Sized> Drop for SpinMutexGuard<'_, T> {
        fn drop(&mut self) {
            unsafe { self.mutex.unlock() }
        }
    }

    unsafe impl<T: ?Sized + Send> Send for SpinMutex<T> {}
    unsafe impl<T: ?Sized + Send> Sync for SpinMutex<T> {}

    unsafe impl<T: ?Sized> Send for SpinMutexGuard<'_, T> where for<'a> &'a mut T: Send {}
    unsafe impl<T: ?Sized> Sync for SpinMutexGuard<'_, T> where for<'a> &'a mut T: Sync {}

    impl<T> SpinMutex<T> {
        pub const fn new(data: T) -> Self {
            Self {
                lock: AtomicBool::new(false),
                data: UnsafeCell::new(data),
            }
        }

        pub fn into_inner(self) -> T {
            self.data.into_inner()
        }

        pub fn as_mut_ptr(&self) -> *mut T {
            self.data.get()
        }
    }

    impl<T: ?Sized> SpinMutex<T> {
        pub fn try_lock(&self) -> bool {
            self.lock
                .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
                .is_ok()
        }
        pub fn try_lock_weak(&self) -> bool {
            self.lock
                .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
                .is_ok()
        }
        pub fn lock(&self) -> SpinMutexGuard<'_, T> {
            while !self.try_lock() {
                core::hint::spin_loop();
            }
            SpinMutexGuard { mutex: self }
        }

        /// # Safety
        /// The caller must be the logical owner of the locked mutex.
        pub unsafe fn unlock(&self) {
            self.lock.store(false, Ordering::Release);
        }
    }
}

pub use spin_mutex::SpinMutex;

mod once {
    use core::{
        cell::{Cell, UnsafeCell},
        mem::{ManuallyDrop, MaybeUninit},
        sync::atomic::{AtomicU8, Ordering},
    };

    const UNINITIALIZED: u8 = 0;
    const INITIALIZING: u8 = 1;
    const INITIALIZED: u8 = 2;
    const PANICKED: u8 = 3;

    pub struct OnceState {
        poisoned: bool,
        state_to_set: Cell<u8>,
    }

    impl OnceState {
        pub fn is_poisoned(&self) -> bool {
            self.poisoned
        }
        pub fn poison(&self) {
            self.state_to_set.set(PANICKED);
        }
    }

    pub struct Once {
        state: AtomicU8,
    }

    const impl Default for Once {
        fn default() -> Self {
            Self {
                state: AtomicU8::new(UNINITIALIZED),
            }
        }
    }

    impl Once {
        pub const fn new() -> Self {
            Self {
                state: AtomicU8::new(UNINITIALIZED),
            }
        }

        pub fn state(&self) -> u8 {
            self.state.load(Ordering::Acquire)
        }

        pub fn set_state(&self, state: u8) {
            self.state.store(state, Ordering::Release);
        }

        pub fn is_completed(&self) -> bool {
            self.state.load(Ordering::Acquire) == INITIALIZED
        }

        pub fn call_once<F>(&self, f: F)
        where
            F: FnOnce(&OnceState),
        {
            if self.is_completed() {
                return;
            }

            self.call_once_slow(false, f);
        }

        pub fn call_once_force<F>(&self, f: F)
        where
            F: FnOnce(&OnceState),
        {
            if self.is_completed() {
                return;
            }

            self.call_once_slow(true, f);
        }

        #[cold]
        pub fn call_once_slow<F>(&self, ignore_poison: bool, f: F)
        where
            F: FnOnce(&OnceState),
        {
            let mut state = self.state.load(Ordering::Acquire);
            loop {
                match state {
                    INITIALIZED => return,
                    PANICKED if !ignore_poison => {
                        panic!("Once instance has previously been poisoned")
                    }
                    PANICKED | UNINITIALIZED => {
                        match self.state.compare_exchange(
                            state,
                            INITIALIZING,
                            Ordering::Acquire,
                            // if we get `Err(INITIALIZED)`, we want to have
                            // acquired what the lock is protecting.
                            Ordering::Acquire,
                        ) {
                            Err(new) => {
                                state = new;
                                continue;
                            }
                            Ok(_) => {
                                // even though we don't have unwinding, for
                                // completeness sake we'll poison the lock if
                                // the closure panics.
                                struct Guard<'a>(&'a AtomicU8);

                                impl Drop for Guard<'_> {
                                    fn drop(&mut self) {
                                        self.0.store(PANICKED, Ordering::Release);
                                    }
                                }

                                let guard = Guard(&self.state);

                                let state = OnceState {
                                    poisoned: state == PANICKED,
                                    state_to_set: Cell::new(INITIALIZED),
                                };

                                f(&state);
                                core::mem::forget(guard);

                                self.state
                                    .store(state.state_to_set.take(), Ordering::Release);
                                return;
                            }
                        }
                    }
                    _ => {
                        assert_eq!(state, INITIALIZING);
                        loop {
                            state = self.state.load(Ordering::Acquire);
                            if state != INITIALIZING {
                                break;
                            }
                            core::hint::spin_loop();
                        }
                    }
                }
            }
        }
    }

    pub struct OnceLock<T> {
        once: Once,
        data: UnsafeCell<MaybeUninit<T>>,
    }

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

    const impl<T> Default for OnceLock<T> {
        fn default() -> Self {
            Self {
                once: Once::new(),
                data: UnsafeCell::new(MaybeUninit::uninit()),
            }
        }
    }

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

        pub const fn from(t: T) -> Self {
            Self {
                once: Once::new(),
                data: UnsafeCell::new(MaybeUninit::new(t)),
            }
        }

        /// # Safety
        /// The caller must ensure that the `OnceLock` has been initialized before calling this method.
        pub unsafe fn get_unchecked(&self) -> &T {
            unsafe { self.data.get().as_ref_unchecked().assume_init_ref() }
        }

        /// # Safety
        /// The caller must ensure that the `OnceLock` has been initialized before calling this method.
        pub unsafe fn get_mut_unchecked(&mut self) -> &mut T {
            unsafe { self.data.get().as_mut_unchecked().assume_init_mut() }
        }

        pub fn get(&self) -> Option<&T> {
            if self.is_completed() {
                Some(unsafe { self.get_unchecked() })
            } else {
                None
            }
        }

        pub fn get_mut(&mut self) -> Option<&mut T> {
            if self.is_completed_mut() {
                Some(unsafe { self.get_mut_unchecked() })
            } else {
                None
            }
        }

        pub fn is_completed(&self) -> bool {
            self.once.is_completed()
        }

        pub fn is_completed_mut(&mut self) -> bool {
            // avoid atomic load when we have mutable access to self
            *self.once.state.get_mut() == INITIALIZED
        }

        pub fn get_spinning(&self) -> Option<&T> {
            loop {
                match self.once.state() {
                    INITIALIZED => return Some(unsafe { self.get_unchecked() }),
                    INITIALIZING => core::hint::spin_loop(),
                    UNINITIALIZED | PANICKED => return None,
                    _ => unreachable!(),
                }
            }
        }

        pub fn initialize<F, E>(&self, f: F) -> Result<(), E>
        where
            F: FnOnce() -> Result<T, E>,
        {
            let mut res: Result<(), E> = Ok(());
            let cell = unsafe { &mut *self.data.get() };

            self.once.call_once_force(|state| match f() {
                Ok(value) => {
                    unsafe { cell.as_mut_ptr().write(value) };
                }
                Err(e) => {
                    state.poison();
                    res = Err(e);
                }
            });

            res
        }

        pub fn get_or_init<F>(&self, f: F) -> &T
        where
            F: FnOnce() -> T,
        {
            match self.get_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
                Ok(value) => value,
                Err(_) => unreachable!(),
            }
        }

        pub fn get_mut_or_init<F>(&mut self, f: F) -> &mut T
        where
            F: FnOnce() -> T,
        {
            match self.get_mut_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
                Ok(value) => value,
                Err(_) => unreachable!(),
            }
        }

        pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
        where
            F: FnOnce() -> Result<T, E>,
        {
            if !self.is_completed() {
                self.initialize(f)?;
            }

            unsafe { Ok(self.get_unchecked()) }
        }

        pub fn get_mut_or_try_init<F, E>(&mut self, f: F) -> Result<&mut T, E>
        where
            F: FnOnce() -> Result<T, E>,
        {
            if self.get_mut().is_none() {
                self.initialize(f)?;
            }

            Ok(unsafe { self.get_mut_unchecked() })
        }

        pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
            let mut val = Some(value);
            let res = self.get_or_init(|| unsafe { val.take().unwrap_unchecked() });
            match val {
                Some(value) => Err((res, value)),
                None => Ok(res),
            }
        }
    }

    impl<T> Drop for OnceLock<T> {
        fn drop(&mut self) {
            if self.is_completed() {
                unsafe { self.data.get().as_mut_unchecked().assume_init_drop() }
            }
        }
    }

    pub struct LazyLock<T, F = fn() -> T> {
        once: Once,
        f: UnsafeCell<ManuallyDrop<F>>,
        t: UnsafeCell<MaybeUninit<T>>,
    }

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

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

        pub fn get(&self) -> Option<&T> {
            if self.once.is_completed() {
                Some(unsafe { (&*self.t.get()).assume_init_ref() })
            } else {
                None
            }
        }

        fn force(this: &Self) -> &T {
            this.once.call_once(|_| {
                // SAFETY: because `call_once` will panic if poisoned, this
                // closure will only be called once.
                let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) };
                let val = f();
                unsafe { (&mut *this.t.get()).write(val) };
            });

            unsafe { (&*this.t.get()).assume_init_ref() }
        }

        fn force_mut(this: &mut Self) -> &mut T {
            this.once.call_once(|_| {
                // SAFETY: because `call_once` will panic if poisoned, this
                // closure will only be called once.
                let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) };
                let val = f();
                unsafe { (&mut *this.t.get()).write(val) };
            });

            unsafe { (&mut *this.t.get()).assume_init_mut() }
        }
    }

    impl<T, F: FnOnce() -> T> core::ops::Deref for LazyLock<T, F> {
        type Target = T;

        fn deref(&self) -> &Self::Target {
            LazyLock::force(self)
        }
    }

    impl<T, F: FnOnce() -> T> core::ops::DerefMut for LazyLock<T, F> {
        fn deref_mut(&mut self) -> &mut Self::Target {
            LazyLock::force_mut(self)
        }
    }

    impl<T, F> Drop for LazyLock<T, F> {
        fn drop(&mut self) {
            match self.once.state() {
                UNINITIALIZED => unsafe {
                    ManuallyDrop::drop(self.f.get_mut());
                },
                INITIALIZED => unsafe {
                    (&mut *self.t.get()).assume_init_drop();
                },
                _ => {}
            }
        }
    }
}

pub use once::{LazyLock, Once, OnceLock};