odem-rs-core 0.1.0

Core components of the ODEM-rs simulation framework
//! The `irc` module provides an intrusive reference-counting smart pointer,
//! `Irc`, along with traits for implementing intrusively counted types.

use super::LeasedMut;
use core::{
	any::{Any, TypeId},
	borrow::Borrow,
	fmt,
	ops::Deref,
	panic::{Location, RefUnwindSafe, UnwindSafe},
	pin::Pin,
	ptr::NonNull,
};
use std::ops::DerefMut;

/* ******************************************************************* Traits */

/// Used to do a cheap reference-to-[Irc] conversion.
pub trait AsIrc<T: ?Sized + IntrusivelyCounted> {
	/// Converts this type into an [owned reference] of the (usually inferred)
	/// input type.
	///
	/// [owned reference]: Irc
	fn as_irc(&self) -> Irc<T>;
}

/// A trait that can be implemented by types that track references using an
/// internal reference counter.
///
/// # Safety
/// The implementation must guarantee that the `IrcBox` is owned by the
/// implementor, i.e. will be dropped together with `Self`. This usually means
/// that the `IrcBox` is embedded in the type implementing this trait.
pub unsafe trait IntrusivelyCounted {
	/// The intrusive type containing the reference count and additional
	/// meta-data.
	type Inner: IrcBoxed;

	/// Returns a reference to the embedded object implementing [IrcBoxed].
	/// It is used by [Irc] to opaquely manipulate the reference counter.
	fn irc_box(&self) -> &IrcBox<Self::Inner>;
}

/// A trait allowing an [`Irc`] to increment and decrement the reference count
/// as well as recycling an unreachable object.
///
/// # Safety
/// The implementor has to ensure that the `acquire` and `release` methods
/// provided by the trait correctly keep track of the inner reference counts
/// in order for the `Irc` derived from them to be sound.
pub unsafe trait IrcBoxed {
	/// Returns the number of `Irc` currently referencing the `IrcBox`.
	fn ref_count(&self) -> usize;

	/// Increments the reference counter.
	fn acquire(&self, _: Private);

	/// Decrements the internal reference counter, returning a function pointer
	/// responsible for recycling if it reaches zero.
	///
	/// The reason for this convoluted mechanism is that the implementing
	/// [`IrcBox`] may recover the original type and reclaim it using an
	/// exclusive reference. This cannot be done in the function directly
	/// because a shared reference to the `IrcBox` exists, precluding the
	/// existence of mutable references to objects that include the `IrcBox`
	/// itself.
	fn release(&self, _: Private) -> Option<fn(NonNull<Self>)>;
}

/* ************************************************************ Irc Structure */

/// The intrusive version of a [Rc] without support for [Weak] pointer. 'Irc'
/// stands for 'Intrusively Reference Counted'.
///
/// [Rc]: std::rc::Rc
/// [Weak]: std::rc::Weak
pub struct Irc<T: ?Sized + IntrusivelyCounted>(NonNull<T>);

impl<T: ?Sized + IntrusivelyCounted> Irc<T> {
	/// Creates a new intrusively counted pointer from a [`Lease`] value.
	///
	/// Because the value is borrowed exclusively, we can be sure that it
	/// currently is the only reference to the value. Because of the invariant
	/// lifetime `'p` associated with the [`Lease`], we can be sure that it
	/// can *never* be borrowed again.
	/// Finally, because the value is [pinned] and `T` implements
	/// [`IntrusivelyCounted`], we can be sure that its inner [`IrcBox`] will
	/// either be dropped or remain valid indefinitely.
	///
	/// [`Lease`]: super::Lease
	/// [pinned]: core::pin
	pub fn new(value: Pin<LeasedMut<'_, T>>) -> Self {
		// increase the reference count to one
		value.irc_box().inner.acquire(Private(()));

		// strip away the lifetime requirements
		Irc(NonNull::from(&mut **unsafe {
			Pin::into_inner_unchecked(value)
		}))

		// dangling pointers are prevented by the Drop impl of `IrcBox`
	}

	/// Unsafely creates a new intrusively counted pointer from a pinned value.
	///
	/// Because the value is borrowed exclusively, we can be sure that it
	/// currently is the only reference to the value. Because the value is
	/// [pinned](core::pin) and `T` implements [`IntrusivelyCounted`], we can be
	/// sure that its inner [`IrcBox`] will either be dropped or remain valid
	/// indefinitely.
	///
	/// # Safety
	/// We cannot be sure that the caller doesn't hold an external reference
	/// that the value is borrowed from, that will become active later. This is
	/// a problem because we allow mutable borrows in `Irc::drop` based on the
	/// reference count and also through `Irc::as_pin_mut`.
	///
	/// The caller is responsible that the pinned value is not duplicated and
	/// used after this method has been called. Keeping a reference - mutable or
	/// not - results in **undefined behavior**.
	pub unsafe fn new_unchecked(value: Pin<&mut T>) -> Self {
		// increase the reference count to one
		value.irc_box().acquire(Private(()));

		// strip away the lifetime requirements
		Irc(NonNull::from(unsafe { Pin::into_inner_unchecked(value) }))

		// dangling pointers are prevented by the Drop impl of `IrcBox`
	}

	/// Returns a pinned reference to the inner value.
	pub fn get_pin(&self) -> Pin<&T> {
		// SAFETY: the pointer was originally pinned, so reconstructing this
		// constraint here is safe
		unsafe { Pin::new_unchecked(self.0.as_ref()) }
	}

	/// Returns a pinned mutable reference to the inner value if there is only
	/// one `Irc` in use right now.
	pub fn get_pin_mut(&mut self) -> Option<Pin<&mut T>> {
		// SAFETY: the constructors originally required a pinned mutable
		// reference that left no reference with the caller; therefore restoring
		// it if we're sure that there is only this `Irc` is safe
		(self.irc_box().ref_count() == 1)
			.then(|| unsafe { Pin::new_unchecked(self.0.as_mut()) })
	}

	/// Strips the outer structure from the `Irc`, revealing the inner
	/// unprotected reference. This operation does not decrease the reference
	/// counter and should be used in tandem with [`from_raw`] to restore the
	/// smart pointer at a later time.
	///
	/// [`from_raw`]: Self::from_raw
	pub const fn into_raw(this: Self) -> NonNull<T> {
		// copy the inner pointer
		let inner = this.0;

		// forget calling drop
		core::mem::forget(this);

		// return the pointer
		inner
	}

	/// Reconstructs the `Irc` from an unprotected reference.
	///
	/// # Safety
	/// The associated function is marked as unsafe because it is the caller's
	/// responsibility to ensure that this reference has originally been the result
	/// of a call to [`into_raw`].
	///
	/// [`into_raw`]: Self::into_raw
	pub const unsafe fn from_raw(inner: NonNull<T>) -> Self {
		Irc(inner)
	}

	/// Allows the (limited) projection of a composite [Irc] into an Irc of one
	/// of its member variables, provided that member variable contains the same
	/// reference counter.
	pub fn map<F, R>(this: Self, f: F) -> Irc<R>
	where
		F: FnOnce(&T) -> &R,
		R: ?Sized + IntrusivelyCounted,
	{
		// store the pointer to the inner IrcBox
		let inner = this.irc_box();

		// perform the conversion
		let value = f(&*this);

		// assert that the IrcBox hasn't changed due to the projection;
		// this should be optimized away by the compiler
		assert!(
			core::ptr::addr_eq(inner, value.irc_box()),
			"expected the mapping to yield an Irc with the same IrcBox"
		);

		// construct the new Irc, taking the provenance from `this`
		let res = Irc(with_provenance(NonNull::from(value), this.0));

		// forget the original Irc
		core::mem::forget(this);

		res
	}

	/// Converts the `Irc` into type `V` if it is of this type, or returns the
	/// old `Irc` if it isn't.
	pub fn downcast<V>(self) -> Result<Irc<V>, Irc<T>>
	where
		T: Any,
		V: 'static + IntrusivelyCounted,
	{
		if (*self).type_id() == TypeId::of::<V>() {
			let res = Irc(self.0.cast::<V>());
			core::mem::forget(self);
			Ok(res)
		} else {
			Err(self)
		}
	}
}

impl<T: ?Sized + IntrusivelyCounted> Deref for Irc<T> {
	type Target = T;

	fn deref(&self) -> &Self::Target {
		unsafe { self.0.as_ref() }
	}
}

impl<T: ?Sized + IntrusivelyCounted> Clone for Irc<T> {
	fn clone(&self) -> Self {
		self.irc_box().acquire(Private(()));
		Irc(self.0)
	}
}

impl<T: ?Sized + IntrusivelyCounted> Borrow<T> for Irc<T> {
	fn borrow(&self) -> &T {
		self
	}
}

impl<T: ?Sized + IntrusivelyCounted> Drop for Irc<T> {
	fn drop(&mut self) {
		if let Some(recycle) = self.irc_box().release(Private(())) {
			recycle(NonNull::from(&self.irc_box().inner));
		}
	}
}

impl<T: ?Sized + IntrusivelyCounted> Unpin for Irc<T> {}

impl<T: ?Sized + IntrusivelyCounted + fmt::Debug> fmt::Debug for Irc<T> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		unsafe { self.0.as_ref() }.fmt(f)
	}
}

impl<T: ?Sized + IntrusivelyCounted + fmt::Display> fmt::Display for Irc<T> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		unsafe { self.0.as_ref() }.fmt(f)
	}
}

// T: Sync is enough for Irc<T> to be Send + Sync since we're moving references
unsafe impl<T: Sync + ?Sized + IntrusivelyCounted> Send for Irc<T> {}

unsafe impl<T: Sync + ?Sized + IntrusivelyCounted> Sync for Irc<T> {}

impl<T: RefUnwindSafe + ?Sized + IntrusivelyCounted> UnwindSafe for Irc<T> {}

/* ****************************************** Marker Type for Private Details */

/// Marker type to enable generating private functions in a public trait
/// interface.
pub struct Private(());

/* ****************************************************** Provenance Transfer */

/// Transfers the provenance from one non-null pointer to another.
///
/// This function assigns the provenance of the `provenance` pointer to the
/// `target` pointer. Provenance refers to the origin or ownership context of a
/// pointer, which is crucial for maintaining memory safety and correctness in
/// low-level operations. See [here] for Rust's strict provenance.
///
/// # Safety
///
/// This function performs raw pointer manipulation and assumes that the layout
/// of fat pointers with metadata remains consistent.
///
/// [here]: https://doc.rust-lang.org/core/ptr/index.html#strict-provenance
fn with_provenance<S, T>(mut target: NonNull<T>, provenance: NonNull<S>) -> NonNull<T>
where
	S: ?Sized,
	T: ?Sized,
{
	// Create a thin pointer by casting `provenance` to `u8` and setting its
	// address to that of `target`. This combines the address of `target` with
	// the provenance of `provenance`.
	let target_thin_ptr = provenance.cast::<u8>().with_addr(target.addr());

	// Obtain a mutable reference to the `target` pointer and cast it to a
	// pointer to `NonNull<u8>`. This allows direct manipulation of the thin
	// pointer portion of the potentially fat `NonNull<T>`.
	let ptr_to_fat_ptr = NonNull::from(&mut target).cast::<NonNull<u8>>();

	// Overwrite the thin pointer part of the fat `target` pointer with the new
	// thin pointer that carries the desired provenance. This operation
	// preserves the original address while updating its provenance.
	//
	// SAFETY: This is safe provided that the layout of fat pointers with
	// metadata does not change.
	// TODO: use with_metadata_of() once stable
	unsafe {
		ptr_to_fat_ptr.write(target_thin_ptr);
	}

	// Return the updated `target` pointer with the new provenance.
	target
}

/* ********************************************************* IrcBox Structure */

/// Type that has to be stored inside a structure in order to enable creating
/// [`Irc`] to it.
#[derive(Debug)]
pub struct IrcBox<T: ?Sized + IrcBoxed> {
	/// The location in the code.
	loc: &'static Location<'static>,
	/// The inner value containing the reference counters.
	inner: T,
}

impl<T: IrcBoxed> IrcBox<T> {
	/// Creates a new `IrcBox` with the inner value.
	#[track_caller]
	pub const fn new(inner: T) -> Self {
		IrcBox::with_location(inner, Location::caller())
	}

	/// Creates a new `IrcBox` with a specific [Location].
	pub const fn with_location(inner: T, loc: &'static Location<'static>) -> Self {
		IrcBox { loc, inner }
	}

	/// Returns [`Location`]-information related to the creation of the `Irc`.
	pub const fn location(this: &Self) -> &'static Location<'static> {
		this.loc
	}
}

impl<T: IrcBoxed + Default> Default for IrcBox<T> {
	#[track_caller]
	fn default() -> Self {
		Self::new(T::default())
	}
}

impl<T: IrcBoxed> Deref for IrcBox<T> {
	type Target = T;

	fn deref(&self) -> &Self::Target {
		&self.inner
	}
}

impl<T: IrcBoxed> DerefMut for IrcBox<T> {
	fn deref_mut(&mut self) -> &mut Self::Target {
		&mut self.inner
	}
}

impl<T: ?Sized + IrcBoxed> Drop for IrcBox<T> {
	fn drop(&mut self) {
		// prime an abort guard in case we have to panic
		let abort_guard = scopeguard::guard(self.loc, |loc| {
			panic!("aborting due to dangling `Irc` created at '{loc}'");
		});

		// ensure that no references point to this instance
		match self.inner.ref_count() {
			// abort aborting
			0 => core::mem::forget(abort_guard),
			// this panic leads to an immediate abort which is necessary
			// because running all the destructors in the stack frames above
			// will perform a ton of illegal memory accesses
			n => panic!(
				"dropping the value created at '{}' leaves {} reference(s) dangling",
				self.loc, n,
			),
		}
	}
}