test: test both implementations of Mutex/OnceLock

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::<Vec<_>>();
-
-		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() {

src/util/mutex_spin.rs

@@ -67,3 +67,40 @@ impl<T> 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::<Vec<_>>();
+
+		for t in threads {
+			t.join().expect("");
+		}
+
+		assert_eq!(*mutex.lock(), 4000);
+	}
+}

src/util/mutex_std.rs

@@ -16,3 +16,40 @@ impl<T> Mutex<T> {
 		}
 	}
 }
+
+#[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::<Vec<_>>();
+
+		for t in threads {
+			t.join().expect("");
+		}
+
+		assert_eq!(*mutex.lock(), 4000);
+	}
+}

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<T> {
 	data: UnsafeCell<MaybeUninit<T>>,
-	#[cfg(not(feature = "std"))]
 	status: core::sync::atomic::AtomicU8,
-	#[cfg(feature = "std")]
-	once: std::sync::Once,
 	phantom: PhantomData<T>
 }
 
 unsafe impl<T: Send> Send for OnceLock<T> {}
 unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
 
-#[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<T> OnceLock<T> {
 	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<F: FnOnce() -> T>(&self, f: F) -> &T {
-		match self.get_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
+		match self.get_or_try_init(|| Ok::<T, Infallible>(f())) {
 			Ok(x) => x,
 			Err(e) => match e {}
 		}
@@ -66,7 +65,7 @@ impl<T> OnceLock<T> {
 			{
 				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<T> OnceLock<T> {
 					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<T> OnceLock<T> {
 					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<T> OnceLock<T> {
 				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<T> OnceLock<T> {
 			}
 		}
 	}
-}
 
-#[cfg(feature = "std")]
-impl<T> OnceLock<T> {
-	pub const fn new() -> Self {
-		Self {
-			data: UnsafeCell::new(MaybeUninit::uninit()),
-			once: std::sync::Once::new(),
-			phantom: PhantomData
-		}
-	}
-
-	#[inline]
-	pub fn get_or_init<F: FnOnce() -> T>(&self, f: F) -> &T {
-		match self.get_or_try_init(|| Ok::<T, core::convert::Infallible>(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<F: FnOnce() -> Result<T, E>, 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<F: FnOnce() -> Result<T, E>, 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<T> OnceLock<T> {
 	pub fn try_insert_with<F: FnOnce() -> T>(&self, inserter: F) -> bool {
 		let mut container = Some(inserter);
 		self.get_or_init(|| (unsafe { container.take().unwrap_unchecked() })());
@@ -166,14 +114,28 @@ impl<T> OnceLock<T> {
 
 impl<T> Drop for OnceLock<T> {
 	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);
+	}
+}

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<T> {
+	data: UnsafeCell<MaybeUninit<T>>,
+	once: std::sync::Once,
+	phantom: PhantomData<T>
+}
+
+unsafe impl<T: Send> Send for OnceLock<T> {}
+unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
+
+impl<T> OnceLock<T> {
+	pub const fn new() -> Self {
+		Self {
+			data: UnsafeCell::new(MaybeUninit::uninit()),
+			once: Once::new(),
+			phantom: PhantomData
+		}
+	}
+
+	#[inline]
+	pub fn get_or_init<F: FnOnce() -> T>(&self, f: F) -> &T {
+		match self.get_or_try_init(|| Ok::<T, Infallible>(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<F: FnOnce() -> Result<T, E>, 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<F: FnOnce() -> Result<T, E>, 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<T> OnceLock<T> {
+	pub fn try_insert_with<F: FnOnce() -> T>(&self, inserter: F) -> bool {
+		let mut container = Some(inserter);
+		self.get_or_init(|| (unsafe { container.take().unwrap_unchecked() })());
+		container.is_none()
+	}
+}
+
+impl<T> Drop for OnceLock<T> {
+	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);
+	}
+}