From 477e92949589bbd09ea5ca51416b15ab5ef6628e Mon Sep 17 00:00:00 2001 From: "Carson M." Date: Mon, 16 Mar 2026 14:25:48 -0500 Subject: [PATCH] test: test both implementations of `Mutex`/`OnceLock` --- src/util/mod.rs | 56 ++++------ src/util/mutex_spin.rs | 37 +++++++ src/util/mutex_std.rs | 37 +++++++ src/util/{once_lock.rs => once_lock_spin.rs} | 108 ++++++------------- src/util/once_lock_std.rs | 95 ++++++++++++++++ 5 files changed, 222 insertions(+), 111 deletions(-) rename src/util/{once_lock.rs => once_lock_spin.rs} (55%) create mode 100644 src/util/once_lock_std.rs diff --git a/src/util/mod.rs b/src/util/mod.rs index b191c0d..a8ed514 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -8,15 +8,25 @@ use core::{ use crate::{Result, memory::Allocator}; -#[cfg(feature = "std")] -#[path = "mutex_std.rs"] -mod mutex; +#[cfg(any(test, not(feature = "std")))] +mod mutex_spin; +#[cfg(any(test, feature = "std"))] +mod mutex_std; +#[cfg(any(test, not(feature = "std")))] +mod once_lock_spin; +#[cfg(any(test, feature = "std"))] +mod once_lock_std; + #[cfg(not(feature = "std"))] -#[path = "mutex_spin.rs"] -mod mutex; +use mutex_spin as mutex; +#[cfg(feature = "std")] +use mutex_std as mutex; +#[cfg(not(feature = "std"))] +use once_lock_spin as once_lock; +#[cfg(feature = "std")] +use once_lock_std as once_lock; mod map; -mod once_lock; mod stack; pub(crate) use self::{ map::MiniMap, @@ -180,40 +190,10 @@ impl<'a, 'p> Drop for AllocatedString<'a, 'p> { #[cfg(test)] #[allow(clippy::unwrap_used)] mod tests { - use alloc::{ffi::CString, sync::Arc}; + use alloc::ffi::CString; use core::ffi::CStr; - use std::thread; - use super::{MiniMap, Mutex, char_p_to_string, run_on_drop, with_cstr, with_cstr_ptr_array}; - - #[test] - fn test_mutex_sanity() { - let mutex = Mutex::new(()); - for _ in 0..4 { - drop(mutex.lock()); - } - } - - #[test] - fn test_mutex_threaded() { - let mutex = Arc::new(Mutex::new(0usize)); - let threads = (0..4) - .map(|_| { - let mutex = Arc::clone(&mutex); - thread::spawn(move || { - for _ in 0..1000 { - *mutex.lock() += 1; - } - }) - }) - .collect::>(); - - for t in threads { - t.join().unwrap(); - } - - assert_eq!(*mutex.lock(), 4000); - } + use super::{MiniMap, char_p_to_string, run_on_drop, with_cstr, with_cstr_ptr_array}; #[test] fn test_run_on_drop() { diff --git a/src/util/mutex_spin.rs b/src/util/mutex_spin.rs index 2d9df74..fabf2da 100644 --- a/src/util/mutex_spin.rs +++ b/src/util/mutex_spin.rs @@ -67,3 +67,40 @@ impl Drop for MutexGuard<'_, T> { self.is_locked.store(false, Ordering::Release); } } + +#[cfg(test)] +mod tests { + use alloc::sync::Arc; + use std::thread; + + use super::Mutex; + + #[test] + fn test_mutex_sanity() { + let mutex = Mutex::new(()); + for _ in 0..4 { + drop(mutex.lock()); + } + } + + #[test] + fn test_mutex_threaded() { + let mutex = Arc::new(Mutex::new(0usize)); + let threads = (0..4) + .map(|_| { + let mutex = Arc::clone(&mutex); + thread::spawn(move || { + for _ in 0..1000 { + *mutex.lock() += 1; + } + }) + }) + .collect::>(); + + for t in threads { + t.join().expect(""); + } + + assert_eq!(*mutex.lock(), 4000); + } +} diff --git a/src/util/mutex_std.rs b/src/util/mutex_std.rs index ffa0791..ada1899 100644 --- a/src/util/mutex_std.rs +++ b/src/util/mutex_std.rs @@ -16,3 +16,40 @@ impl Mutex { } } } + +#[cfg(test)] +mod tests { + use alloc::sync::Arc; + use std::thread; + + use super::Mutex; + + #[test] + fn test_mutex_sanity() { + let mutex = Mutex::new(()); + for _ in 0..4 { + drop(mutex.lock()); + } + } + + #[test] + fn test_mutex_threaded() { + let mutex = Arc::new(Mutex::new(0usize)); + let threads = (0..4) + .map(|_| { + let mutex = Arc::clone(&mutex); + thread::spawn(move || { + for _ in 0..1000 { + *mutex.lock() += 1; + } + }) + }) + .collect::>(); + + for t in threads { + t.join().expect(""); + } + + assert_eq!(*mutex.lock(), 4000); + } +} diff --git a/src/util/once_lock.rs b/src/util/once_lock_spin.rs similarity index 55% rename from src/util/once_lock.rs rename to src/util/once_lock_spin.rs index b11aaeb..8503696 100644 --- a/src/util/once_lock.rs +++ b/src/util/once_lock_spin.rs @@ -1,39 +1,38 @@ -#[cfg(not(feature = "std"))] -use core::sync::atomic::Ordering; -use core::{cell::UnsafeCell, marker::PhantomData, mem::MaybeUninit}; +use core::{ + cell::UnsafeCell, + convert::Infallible, + hint::spin_loop, + marker::PhantomData, + mem::{MaybeUninit, forget}, + ptr, + sync::atomic::{AtomicU8, Ordering} +}; pub(crate) struct OnceLock { data: UnsafeCell>, - #[cfg(not(feature = "std"))] status: core::sync::atomic::AtomicU8, - #[cfg(feature = "std")] - once: std::sync::Once, phantom: PhantomData } unsafe impl Send for OnceLock {} unsafe impl Sync for OnceLock {} -#[cfg(not(feature = "std"))] const STATUS_UNINITIALIZED: u8 = 0; -#[cfg(not(feature = "std"))] const STATUS_RUNNING: u8 = 1; -#[cfg(not(feature = "std"))] const STATUS_INITIALIZED: u8 = 2; -#[cfg(not(feature = "std"))] impl OnceLock { pub const fn new() -> Self { Self { data: UnsafeCell::new(MaybeUninit::uninit()), - status: core::sync::atomic::AtomicU8::new(STATUS_UNINITIALIZED), + status: AtomicU8::new(STATUS_UNINITIALIZED), phantom: PhantomData } } #[inline] pub fn get_or_init T>(&self, f: F) -> &T { - match self.get_or_try_init(|| Ok::(f())) { + match self.get_or_try_init(|| Ok::(f())) { Ok(x) => x, Err(e) => match e {} } @@ -66,7 +65,7 @@ impl OnceLock { { Ok(_) => { struct SetStatusOnPanic<'a> { - status: &'a core::sync::atomic::AtomicU8 + status: &'a AtomicU8 } impl Drop for SetStatusOnPanic<'_> { fn drop(&mut self) { @@ -78,7 +77,7 @@ impl OnceLock { let val = match f() { Ok(val) => val, Err(err) => { - core::mem::forget(panic_catcher); + forget(panic_catcher); self.status.store(STATUS_UNINITIALIZED, Ordering::Release); return Err(err); } @@ -86,7 +85,7 @@ impl OnceLock { unsafe { (*self.data.get()).as_mut_ptr().write(val); }; - core::mem::forget(panic_catcher); + forget(panic_catcher); self.status.store(STATUS_INITIALIZED, Ordering::Release); @@ -95,7 +94,7 @@ impl OnceLock { Err(STATUS_INITIALIZED) => return Ok(unsafe { self.get_unchecked() }), Err(STATUS_RUNNING) => loop { match self.status.load(Ordering::Acquire) { - STATUS_RUNNING => core::hint::spin_loop(), + STATUS_RUNNING => spin_loop(), STATUS_INITIALIZED => return Ok(unsafe { self.get_unchecked() }), // STATUS_UNINITIALIZED - running thread failed, time for us to step in _ => continue 'a @@ -105,58 +104,7 @@ impl OnceLock { } } } -} -#[cfg(feature = "std")] -impl OnceLock { - pub const fn new() -> Self { - Self { - data: UnsafeCell::new(MaybeUninit::uninit()), - once: std::sync::Once::new(), - phantom: PhantomData - } - } - - #[inline] - pub fn get_or_init T>(&self, f: F) -> &T { - match self.get_or_try_init(|| Ok::(f())) { - Ok(x) => x, - Err(e) => match e {} - } - } - - #[inline] - pub fn get(&self) -> Option<&T> { - if self.once.is_completed() { Some(unsafe { self.get_unchecked() }) } else { None } - } - - #[inline] - pub unsafe fn get_unchecked(&self) -> &T { - unsafe { &*(*self.data.get()).as_ptr() } - } - - #[inline] - pub fn get_or_try_init Result, E>(&self, f: F) -> Result<&T, E> { - if let Some(value) = self.get() { Ok(value) } else { self.try_init_inner(f) } - } - - #[cold] - fn try_init_inner Result, E>(&self, f: F) -> Result<&T, E> { - let mut res: Result<(), E> = Ok(()); - let slot = &self.data; - self.once.call_once_force(|_| match f() { - Ok(value) => unsafe { - (*slot.get()).write(value); - }, - Err(e) => { - res = Err(e); - } - }); - res.map(|_| unsafe { self.get_unchecked() }) - } -} - -impl OnceLock { pub fn try_insert_with T>(&self, inserter: F) -> bool { let mut container = Some(inserter); self.get_or_init(|| (unsafe { container.take().unwrap_unchecked() })()); @@ -166,14 +114,28 @@ impl OnceLock { impl Drop for OnceLock { fn drop(&mut self) { - #[cfg(not(feature = "std"))] - let status = *self.status.get_mut() == STATUS_INITIALIZED; - #[cfg(feature = "std")] - let status = self.once.is_completed(); - if status { + if *self.status.get_mut() == STATUS_INITIALIZED { unsafe { - core::ptr::drop_in_place((*self.data.get()).as_mut_ptr()); + ptr::drop_in_place((*self.data.get()).as_mut_ptr()); } } } } + +#[cfg(test)] +mod tests { + use super::OnceLock; + + #[test] + fn test_once() { + let once = OnceLock::new(); + let mut called = 0; + once.get_or_init(|| called += 1); + assert_eq!(called, 1); + once.get_or_init(|| called += 1); + assert_eq!(called, 1); + + assert!(!once.try_insert_with(|| called += 1)); + assert_eq!(called, 1); + } +} diff --git a/src/util/once_lock_std.rs b/src/util/once_lock_std.rs new file mode 100644 index 0000000..1f16541 --- /dev/null +++ b/src/util/once_lock_std.rs @@ -0,0 +1,95 @@ +use core::{cell::UnsafeCell, convert::Infallible, marker::PhantomData, mem::MaybeUninit, ptr}; +use std::sync::Once; + +pub(crate) struct OnceLock { + data: UnsafeCell>, + once: std::sync::Once, + phantom: PhantomData +} + +unsafe impl Send for OnceLock {} +unsafe impl Sync for OnceLock {} + +impl OnceLock { + pub const fn new() -> Self { + Self { + data: UnsafeCell::new(MaybeUninit::uninit()), + once: Once::new(), + phantom: PhantomData + } + } + + #[inline] + pub fn get_or_init T>(&self, f: F) -> &T { + match self.get_or_try_init(|| Ok::(f())) { + Ok(x) => x, + Err(e) => match e {} + } + } + + #[inline] + pub fn get(&self) -> Option<&T> { + if self.once.is_completed() { Some(unsafe { self.get_unchecked() }) } else { None } + } + + #[inline] + pub unsafe fn get_unchecked(&self) -> &T { + unsafe { &*(*self.data.get()).as_ptr() } + } + + #[inline] + pub fn get_or_try_init Result, E>(&self, f: F) -> Result<&T, E> { + if let Some(value) = self.get() { Ok(value) } else { self.try_init_inner(f) } + } + + #[cold] + fn try_init_inner Result, E>(&self, f: F) -> Result<&T, E> { + let mut res: Result<(), E> = Ok(()); + let slot = &self.data; + self.once.call_once_force(|_| match f() { + Ok(value) => unsafe { + (*slot.get()).write(value); + }, + Err(e) => { + res = Err(e); + } + }); + res.map(|_| unsafe { self.get_unchecked() }) + } +} + +impl OnceLock { + pub fn try_insert_with T>(&self, inserter: F) -> bool { + let mut container = Some(inserter); + self.get_or_init(|| (unsafe { container.take().unwrap_unchecked() })()); + container.is_none() + } +} + +impl Drop for OnceLock { + fn drop(&mut self) { + if self.once.is_completed() { + unsafe { + ptr::drop_in_place((*self.data.get()).as_mut_ptr()); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::OnceLock; + + #[test] + fn test_once() { + let once = OnceLock::new(); + let mut called = 0; + once.get_or_init(|| called += 1); + assert_eq!(called, 1); + once.get_or_init(|| called += 1); + assert_eq!(called, 1); + + assert!(!once.try_insert_with(|| called += 1)); + assert_eq!(called, 1); + } +}