From ad8a8e073c07dc936030c31227f2925248367813 Mon Sep 17 00:00:00 2001 From: janis Date: Fri, 17 Jul 2026 16:24:09 +0200 Subject: [PATCH] sync --- kernel/src/lib.rs | 1 + kernel/src/sync/mod.rs | 369 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 370 insertions(+) create mode 100644 kernel/src/sync/mod.rs diff --git a/kernel/src/lib.rs b/kernel/src/lib.rs index 7f32173..2f3cee5 100644 --- a/kernel/src/lib.rs +++ b/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 diff --git a/kernel/src/sync/mod.rs b/kernel/src/sync/mod.rs new file mode 100644 index 0000000..80e24b2 --- /dev/null +++ b/kernel/src/sync/mod.rs @@ -0,0 +1,369 @@ +mod spin_mutex { + use core::{ + cell::UnsafeCell, + sync::atomic::{AtomicBool, Ordering}, + }; + + pub struct SpinMutex { + lock: AtomicBool, + data: UnsafeCell, + } + + pub struct SpinMutexGuard<'a, T: ?Sized + 'a> { + mutex: &'a SpinMutex, + } + + impl core::ops::Deref for SpinMutexGuard<'_, T> { + type Target = T; + + fn deref(&self) -> &Self::Target { + unsafe { &*self.mutex.data.get() } + } + } + + impl core::ops::DerefMut for SpinMutexGuard<'_, T> { + fn deref_mut(&mut self) -> &mut Self::Target { + unsafe { &mut *self.mutex.data.get() } + } + } + + impl Drop for SpinMutexGuard<'_, T> { + fn drop(&mut self) { + unsafe { self.mutex.unlock() } + } + } + + unsafe impl Send for SpinMutex {} + unsafe impl Sync for SpinMutex {} + + unsafe impl Send for SpinMutexGuard<'_, T> where for<'a> &'a mut T: Send {} + unsafe impl Sync for SpinMutexGuard<'_, T> where for<'a> &'a mut T: Sync {} + + impl SpinMutex { + 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 SpinMutex { + 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 { + state: AtomicU8, + data: UnsafeCell>, + } + + unsafe impl Send for Once {} + unsafe impl Sync for Once {} + + const impl Default for Once { + 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 { + 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 Once { + 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(&self, f: F) -> &T + where + F: FnOnce() -> T, + { + if self.get().is_none() { + self.get_or_try_init(|| Ok::(f())); + } + + unsafe { self.get_unchecked() } + } + + pub fn get_or_try_init(&self, f: F) -> Result<&T, E> + where + F: FnOnce() -> Result, + { + 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(&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(&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 Drop for Once { + fn drop(&mut self) { + if self.is_completed() { + unsafe { self.data.get().as_mut_unchecked().assume_init_drop() } + } + } + } + + pub struct LazyLock T> { + once: Once, + init: UnsafeCell>, + } + + unsafe impl Sync for LazyLock where Once: Sync {} + + impl LazyLock + 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> core::ops::Deref for LazyLock { + type Target = T; + + fn deref(&self) -> &Self::Target { + LazyLock::force(self) + } + } + + impl T> core::ops::DerefMut for LazyLock { + fn deref_mut(&mut self) -> &mut Self::Target { + LazyLock::force_mut(self) + } + } + + impl Drop for LazyLock { + fn drop(&mut self) { + if !self.once.is_completed() { + unsafe { + ManuallyDrop::drop(self.init.get_mut()); + } + } + } + } +} + +pub use once::{LazyLock, Once};