sync

kernel/src/lib.rs

@@ -2,6 +2,7 @@
 #![feature(const_trait_impl, const_default)]
 
 pub mod limine;
+pub mod sync;
 pub mod x86_64;
 
 /// # Safety

kernel/src/sync/mod.rs

@@ -0,0 +1,369 @@
+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::UnsafeCell,
+        mem::{ManuallyDrop, MaybeUninit},
+        sync::atomic::{AtomicU8, Ordering},
+    };
+
+    pub struct Once<T> {
+        state: AtomicU8,
+        data: UnsafeCell<MaybeUninit<T>>,
+    }
+
+    unsafe impl<T: Send> Send for Once<T> {}
+    unsafe impl<T: Send + Sync> Sync for Once<T> {}
+
+    const impl<T> Default for Once<T> {
+        fn default() -> Self {
+            Self {
+                state: AtomicU8::new(0),
+                data: UnsafeCell::new(MaybeUninit::uninit()),
+            }
+        }
+    }
+
+    #[repr(u8)]
+    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+    enum OnceState {
+        Uninitialized = 0,
+        Initializing = 1,
+        Initialized = 2,
+    }
+    impl OnceState {
+        fn from_u8(value: u8) -> Option<Self> {
+            match value {
+                0 => Some(Self::Uninitialized),
+                1 => Some(Self::Initializing),
+                2 => Some(Self::Initialized),
+                _ => None,
+            }
+        }
+        unsafe fn from_u8_unchecked(value: u8) -> Self {
+            unsafe { core::mem::transmute(value) }
+        }
+        const fn into_u8(self) -> u8 {
+            self as u8
+        }
+    }
+
+    impl<T> Once<T> {
+        pub const fn new() -> Self {
+            Self {
+                state: AtomicU8::new(OnceState::Uninitialized.into_u8()),
+                data: UnsafeCell::new(MaybeUninit::uninit()),
+            }
+        }
+
+        pub const fn from(t: T) -> Self {
+            Self {
+                state: AtomicU8::new(OnceState::Initialized.into_u8()),
+                data: UnsafeCell::new(MaybeUninit::new(t)),
+            }
+        }
+
+        pub unsafe fn get_unchecked(&self) -> &T {
+            unsafe { self.data.get().as_ref_unchecked().assume_init_ref() }
+        }
+        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.state(Ordering::Acquire) == OnceState::Initialized {
+                Some(unsafe { self.get_unchecked() })
+            } else {
+                None
+            }
+        }
+
+        pub fn get_mut(&mut self) -> Option<&mut T> {
+            if self.state(Ordering::Acquire) == OnceState::Initialized {
+                Some(unsafe { self.get_mut_unchecked() })
+            } else {
+                None
+            }
+        }
+
+        pub fn is_completed(&self) -> bool {
+            self.state(Ordering::Acquire) == OnceState::Initialized
+        }
+
+        fn get_spinning(&self) -> Option<&T> {
+            loop {
+                match self.state(Ordering::Acquire) {
+                    OnceState::Uninitialized => return None,
+                    OnceState::Initializing => core::hint::spin_loop(),
+                    OnceState::Initialized => return Some(unsafe { self.get_unchecked() }),
+                }
+            }
+        }
+
+        fn should_try_init(&self) -> bool {
+            match unsafe {
+                self.state
+                    .compare_exchange(
+                        OnceState::Uninitialized.into_u8(),
+                        OnceState::Initializing.into_u8(),
+                        Ordering::Acquire,
+                        Ordering::Relaxed,
+                    )
+                    .map(|v| OnceState::from_u8_unchecked(v))
+                    .map_err(|v| OnceState::from_u8_unchecked(v))
+            } {
+                Ok(_) => true,
+                Err(OnceState::Initializing) => self.get_spinning().is_none(),
+                Err(OnceState::Initialized) => false,
+                Err(_) => unreachable!(),
+            }
+        }
+
+        pub fn get_or_init<F>(&self, f: F) -> &T
+        where
+            F: FnOnce() -> T,
+        {
+            if self.get().is_none() {
+                self.get_or_try_init(|| Ok::<T, core::convert::Infallible>(f()));
+            }
+
+            unsafe { self.get_unchecked() }
+        }
+
+        pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            if !self.should_try_init() {
+                return Ok(unsafe { self.get_unchecked() });
+            }
+
+            struct PanicGuard<'a>(&'a AtomicU8);
+
+            impl Drop for PanicGuard<'_> {
+                fn drop(&mut self) {
+                    self.0
+                        .store(OnceState::Uninitialized.into_u8(), Ordering::Release);
+                }
+            }
+
+            let guard = PanicGuard(&self.state);
+
+            match f() {
+                Ok(value) => {
+                    unsafe { self.data.get().as_mut_unchecked().write(value) };
+
+                    core::mem::forget(guard);
+
+                    self.state
+                        .store(OnceState::Initialized.into_u8(), Ordering::Release);
+
+                    Ok(unsafe { self.get_unchecked() })
+                }
+                Err(e) => {
+                    core::mem::forget(guard);
+                    self.state
+                        .store(OnceState::Uninitialized.into_u8(), Ordering::Release);
+
+                    Err(e)
+                }
+            }
+        }
+
+        pub fn get_mut_or_init<F>(&mut self, f: F) -> &mut T
+        where
+            F: FnOnce() -> T,
+        {
+            if self.get_mut().is_none() {
+                self.init_once(f);
+            }
+
+            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),
+            }
+        }
+
+        fn init_once<F>(&self, f: F) -> &T
+        where
+            F: FnOnce() -> T,
+        {
+            if self.should_try_init() {
+                let value = f();
+                unsafe { self.data.get().as_mut_unchecked().write(value) };
+                self.state
+                    .store(OnceState::Initialized.into_u8(), Ordering::Release);
+            }
+            unsafe { self.get_unchecked() }
+        }
+
+        fn state(&self, ordering: Ordering) -> OnceState {
+            unsafe { OnceState::from_u8_unchecked(self.state.load(ordering)) }
+        }
+    }
+
+    impl<T> Drop for Once<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<T>,
+        init: UnsafeCell<ManuallyDrop<F>>,
+    }
+
+    unsafe impl<T, F: Send> Sync for LazyLock<T, F> where Once<T>: Sync {}
+
+    impl<T, F> LazyLock<T, F>
+    where
+        F: FnOnce() -> T,
+    {
+        pub const fn new(f: F) -> Self {
+            Self {
+                once: Once::new(),
+                init: UnsafeCell::new(ManuallyDrop::new(f)),
+            }
+        }
+
+        pub fn get(&self) -> Option<&T> {
+            self.once.get()
+        }
+
+        fn force(this: &Self) -> &T {
+            // SAFETY: `once` cannot recycle from `Initializing` to `Uninitialized`, so the closure will only be called once.
+            this.once.get_or_init(|| unsafe {
+                let init = &mut *this.init.get();
+                let f = ManuallyDrop::take(init);
+                f()
+            })
+        }
+
+        fn force_mut(this: &mut Self) -> &mut T {
+            // SAFETY: `once` cannot recycle from `Initializing` to `Uninitialized`, so the closure will only be called once.
+            this.once.get_mut_or_init(|| unsafe {
+                let init = &mut *this.init.get();
+                let f = ManuallyDrop::take(init);
+                f()
+            })
+        }
+    }
+
+    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) {
+            if !self.once.is_completed() {
+                unsafe {
+                    ManuallyDrop::drop(self.init.get_mut());
+                }
+            }
+        }
+    }
+}
+
+pub use once::{LazyLock, Once};