hybrid-rc 0.1.0

//! Helper module for atomically stored thread identification

use core::convert::TryInto;
use core::num::NonZeroUsize;
use core::sync::atomic::{AtomicUsize, Ordering};

thread_local! {
	/// Zero-sized thread-local variable to differentiate threads.
	static THREAD_MARKER: () = ();
}

/// A unique identifier for a running thread.
///
/// Uniqueness is guaranteed between running threads. However, the ids of dead
/// threads may be reused.
///
/// There is a chance that this implementation can be replaced by [`std::thread::ThreadId`]
/// when [`as_u64()`] is stabilized.
///
/// **Note:** The current (non platform specific) implementation uses the address of a
/// thread local static variable for thread identification.
///
/// [`as_u64()`]: std::thread::ThreadId::as_u64
#[derive(PartialEq, Eq, Debug, Clone, Copy)]
#[repr(transparent)]
pub(crate) struct ThreadId(NonZeroUsize);

/// An [`Option`]`<`[`ThreadId`]`>` which can be safely shared between threads.
///
/// **Note:** Currently implemented as a wrapper around [`AtomicUsize`].
#[derive(Debug)]
#[repr(transparent)]
pub(crate) struct AtomicOptionThreadId(core::sync::atomic::AtomicUsize);

impl ThreadId {
	/// Creates a new `ThreadId` for the given raw id.
	#[inline(always)]
	const fn new(value: NonZeroUsize) -> Self {
		Self(value)
	}

	/// Gets the id for the thread that invokes it.
	pub fn current_thread() -> Self {
		Self::new(
			THREAD_MARKER
				.try_with(|x| x as *const _ as usize)
				.expect("the thread's local data has already been destroyed")
				.try_into()
				.expect("thread id should never be zero"),
		)
	}
}

/// Converts the internal representation in [`AtomicOptionThreadId`] into
/// `Some(ThreadId)` or `None`.
///
/// This should only be a type-level conversion and a noop at runtime.
#[inline(always)]
fn wrap(value: usize) -> Option<ThreadId> {
	match value {
		0 => None,
		n => Some(ThreadId::new(n.try_into().unwrap())),
	}
}

/// Converts an `Option<ThreadId>` into the internal representation for
/// [`AtomicOptionThreadId`].
///
/// This should only be a type-level conversion and a noop at runtime.
#[inline(always)]
const fn unwrap(value: Option<ThreadId>) -> usize {
	match value {
		None => 0,
		Some(id) => id.0.get(),
	}
}

impl AtomicOptionThreadId {
	/// Creates a new `AtomicOptionThreadId`.
	#[inline]
	pub const fn new(id: Option<ThreadId>) -> Self {
		Self(AtomicUsize::new(unwrap(id)))
	}

	/// Loads a value from the atomic.
	///
	/// The [`Ordering`] may only be `SeqCst`, `Acquire` or `Relaxed`.
	///
	/// # Panics
	///
	/// Panics if `order` is `Release` or `AcqRel`.
	#[inline]
	pub fn load(&self, order: Ordering) -> Option<ThreadId> {
		wrap(self.0.load(order))
	}

	/// Stores a value into the atomic.
	///
	/// The [`Ordering`] may only be `SeqCst`, `Release` or `Relaxed`.
	///
	/// # Panics
	///
	/// Panics if `order` is `Acquire` or `AcqRel`.
	#[inline]
	pub fn store(&self, val: Option<ThreadId>, order: Ordering) {
		self.0.store(unwrap(val), order);
	}

	/// Stores `new` into the atomic iff `current` is the currently stored value.
	///
	/// The success ordering may be any [`Ordering`], but `failure` may only be
	/// `SeqCst`, `Release` or `Relaxed` and must be equivalent to or weaker than
	/// `success`.
	///
	/// # Panics
	///
	/// Panics if `failure` is `Acquire`, `AcqRel` or a stronger ordering than
	/// `success`.
	#[inline]
	pub fn compare_exchange(
		&self,
		current: Option<ThreadId>,
		new: Option<ThreadId>,
		success: Ordering,
		failure: Ordering,
	) -> Result<Option<ThreadId>, Option<ThreadId>> {
		self.0
			.compare_exchange(unwrap(current), unwrap(new), success, failure)
			.map(wrap)
			.map_err(wrap)
	}
}

impl Default for AtomicOptionThreadId {
	/// Creates an `AtomicOptionThreadId` initialized to [`None`].
	#[inline]
	fn default() -> Self {
		Self::new(None)
	}
}

impl From<ThreadId> for AtomicOptionThreadId {
	/// Converts a [`ThreadId`] into an `AtomicOptionThreadId`, wrapping it in [`Some`].
	#[inline]
	fn from(id: ThreadId) -> Self {
		Self::new(Some(id))
	}
}

impl From<Option<ThreadId>> for AtomicOptionThreadId {
	/// Converts an [`Option`]`<`[`ThreadId`]`>` into an `AtomicOptionThreadId`.
	#[inline]
	fn from(id: Option<ThreadId>) -> Self {
		Self::new(id)
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use std::thread;

	/// Tests if the thread id stays the same on the same thread.
	#[test]
	fn test_eq() {
		let a = ThreadId::current_thread();
		let b = ThreadId::current_thread();
		assert_eq!(a, b);
	}

	/// Tests if the thread id of two different threads differ.
	#[test]
	fn test_ne() {
		let a = ThreadId::current_thread();
		let b = thread::spawn(move || ThreadId::current_thread())
			.join()
			.unwrap();
		assert_ne!(a, b);
	}
}