odem_rs_core/ptr/

irc.rs

//! 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,
	cell::Cell,
	fmt,
	ops::{Deref, DerefMut},
	panic::{Location, RefUnwindSafe, UnwindSafe},
	pin::Pin,
	ptr::NonNull,
	sync::atomic::{AtomicUsize, Ordering},
};

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

/// Used to do an inexpensive 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 {
	/// 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<dyn IrcBoxed>;
}

/// 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.
	fn release(&self, _: Private);

	/// Returns a function pointer that is responsible for recycling if the
	/// reference counter 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.
	#[inline(always)]
	fn reclaim(&self, _: Private) -> Option<fn(NonNull<dyn IrcBoxed>)> {
		None
	}
}

/* ************************************************************ 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().count.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()) })
	}

	/// Returns the inner (raw) pointer of the `Irc`.
	pub fn as_raw(this: &Self) -> NonNull<T> {
		this.0
	}

	/// 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 fn into_raw(this: Self) -> NonNull<T> {
		// copy the inner pointer
		let inner = this.0;

		// forget calling the destructor
		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 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) {
		let irc_box = self.irc_box();
		irc_box.release(Private(()));

		if irc_box.ref_count() == 0 {
			let Some(reclaim) = irc_box.reclaim(Private(())) else {
				return;
			};

			let irc_box = with_provenance(NonNull::from(&irc_box.count), self.0);
			reclaim(irc_box);
		}
	}
}

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: ?Sized, T: ?Sized>(
	mut target: NonNull<T>,
	provenance: NonNull<S>,
) -> NonNull<T> {
	// 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());

	// Get 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 to allow creating [`Irc`] to
/// it.
#[derive(Debug)]
pub struct IrcBox<T: ?Sized + IrcBoxed = Cell<usize>> {
	/// The location in the code.
	loc: &'static Location<'static>,
	/// The inner value containing the reference counters.
	count: T,
}

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

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

	/// 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: ?Sized + IrcBoxed> Deref for IrcBox<T> {
	type Target = T;

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

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

impl<T: ?Sized + IrcBoxed> Drop for IrcBox<T> {
	fn drop(&mut self) {
		// declare a function that cannot unwind when called
		extern "C" fn abort(loc: &Location<'static>, n: usize) -> ! {
			panic!(
				"dropping the value created at '{}' leaves {} reference(s) dangling",
				loc, n,
			)
		}

		// ensure that no references point to this instance
		match self.count.ref_count() {
			0 => {}
			// 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 => abort(self.loc, n),
		}
	}
}

/* ****************************************************** Non-Thread-Safe Box */

// SAFETY: The reference counter is changed and returned accordingly.
unsafe impl IrcBoxed for Cell<usize> {
	#[inline(always)]
	fn ref_count(&self) -> usize {
		self.get()
	}

	#[inline(always)]
	fn acquire(&self, _: Private) {
		self.set(self.get() + 1);
	}

	#[inline(always)]
	fn release(&self, _: Private) {
		self.set(self.get() - 1);
	}
}

/* ********************************************************** Thread-Safe Box */

// SAFETY: The reference counter is changed and returned accordingly.
unsafe impl IrcBoxed for AtomicUsize {
	#[inline(always)]
	fn ref_count(&self) -> usize {
		self.load(Ordering::Relaxed)
	}

	#[inline(always)]
	fn acquire(&self, _: Private) {
		self.fetch_add(1, Ordering::Relaxed);
	}

	#[inline(always)]
	fn release(&self, _: Private) {
		self.fetch_sub(1, Ordering::Relaxed);
	}
}

/* *************************************************************** Miri Tests */

// #[cfg(all(test, miri))]
#[cfg(test)]
mod tests {
	use super::*;
	use crate::ptr::Lease;
	use core::pin::pin;

	/// Tests the correct transfer of provenance during mapping operations, so
	/// that gaining mutable access to the inner `IrcBox` is still defined when
	/// the innermost `Irc` projection is dropped.
	#[test]
	fn provenance_transfer() {
		#[derive(Default)]
		struct Outer {
			inner: Inner,
		}

		#[derive(Default)]
		struct Inner {
			irc_box: IrcBox<InnerBox>,
		}

		#[derive(Default)]
		struct InnerBox {
			rc: Cell<usize>,
			term: bool,
		}

		unsafe impl IrcBoxed for InnerBox {
			fn ref_count(&self) -> usize {
				self.rc.get()
			}

			fn acquire(&self, _: Private) {
				self.rc.set(self.rc.get() + 1);
			}

			fn release(&self, _: Private) {
				self.rc.set(self.rc.get() - 1);
			}

			fn reclaim(&self, _: Private) -> Option<fn(NonNull<dyn IrcBoxed>)> {
				// Return the function that performs the actual reclamation
				Some(|this| unsafe {
					// This mutable access is what requires correct provenance.
					// If Irc::map didn't transfer provenance correctly, Miri
					// might flag this as UB.
					this.cast::<Self>().as_mut().term = true;
				})
			}
		}

		unsafe impl IntrusivelyCounted for Outer {
			fn irc_box(&self) -> &IrcBox<dyn IrcBoxed> {
				&self.inner.irc_box
			}
		}

		unsafe impl IntrusivelyCounted for Inner {
			fn irc_box(&self) -> &IrcBox<dyn IrcBoxed> {
				&self.irc_box
			}
		}

		// 1. Create the data on the stack.
		let outer: Pin<&mut Lease<'_, Outer>> = pin!(Lease::new(Outer::default()));

		// 2. Create an Irc pointer to the outer structure.
		let irc1: Irc<Outer> = Irc::new(outer);

		// 3. Project the Irc<Outer> to an Irc<Inner>.
		let irc2: Irc<Inner> = Irc::map(irc1, |outer| &outer.inner);
		assert_eq!(irc2.irc_box.ref_count(), 1);

		// 4. Drop the Irc<Inner>. This should be the last reference,
		//    triggering the release and reclamation logic.
		drop(irc2);
	}

	/// Tests the correctness of provenance-transfer for an inner `IrcBox`
	/// accessing the outer structure upon reclaim.
	#[test]
	fn self_reference() {
		#[derive(Default)]
		struct Outer {
			inner: Inner,
			done: bool,
		}

		#[derive(Default)]
		struct Inner {
			irc_box: IrcBox<InnerBox>,
		}

		#[derive(Default)]
		struct InnerBox {
			rc: Cell<usize>,
			outer_ref: Cell<Option<NonNull<Outer>>>,
		}

		unsafe impl IrcBoxed for InnerBox {
			fn ref_count(&self) -> usize {
				self.rc.get()
			}

			fn acquire(&self, _: Private) {
				self.rc.set(self.rc.get() + 1);
			}

			fn release(&self, _: Private) {
				self.rc.set(self.rc.get() - 1);
			}

			fn reclaim(&self, _: Private) -> Option<fn(NonNull<dyn IrcBoxed>)> {
				Some(|this| unsafe {
					let mut outer_ref = this
						.cast::<Self>()
						.as_ref()
						.outer_ref
						.get()
						.expect("self-reference initialized");
					outer_ref.as_mut().done = true;
				})
			}
		}

		unsafe impl IntrusivelyCounted for Outer {
			fn irc_box(&self) -> &IrcBox<dyn IrcBoxed> {
				&self.inner.irc_box
			}
		}

		unsafe impl IntrusivelyCounted for Inner {
			fn irc_box(&self) -> &IrcBox<dyn IrcBoxed> {
				&self.irc_box
			}
		}

		// Create a composite structure with a member that includes an IrcBox.
		let outer = pin!(Lease::new(Outer::default()));

		// Create an Irc-pointer to that composite structure.
		let irc1 = Irc::new(outer);

		// Derive a reference to the outmost instance and store it in the IrcBox.
		let outer_ref = Irc::as_raw(&irc1);
		irc1.inner.irc_box.outer_ref.set(Some(outer_ref));

		// Project the Irc onto the inner one, dropping the outer reference.
		let irc2 = Irc::map(irc1, |inner| &inner.inner);

		// Now drop the inner reference, leading to mut access of Outer.
		drop(irc2);
	}

	/// Tests the correctness of provenance transfer for access to an `Outer`
	/// structure upon reclaim.
	#[test]
	fn self_reference2() {
		#[derive(Default)]
		struct Outer {
			inner: Inner,
			done: bool,
		}

		#[derive(Default)]
		struct Inner {
			irc_box: IrcBox<InnerBox>,
		}

		#[derive(Default)]
		struct InnerBox {
			ref_count: Cell<usize>,
			outer_ref: Cell<Option<NonNull<Outer>>>,
		}

		unsafe impl IrcBoxed for InnerBox {
			fn ref_count(&self) -> usize {
				self.ref_count.get()
			}

			fn acquire(&self, _: Private) {
				self.ref_count.set(self.ref_count.get() + 1);
			}

			fn release(&self, _: Private) {
				self.ref_count.set(self.ref_count.get() - 1);
			}

			fn reclaim(&self, _: Private) -> Option<fn(NonNull<dyn IrcBoxed>)> {
				Some(|this| unsafe {
					let mut outer_ref = this
						.cast::<Self>()
						.as_ref()
						.outer_ref
						.get()
						.expect("self-reference initialized");
					outer_ref.as_mut().done = true;
				})
			}
		}

		unsafe impl IntrusivelyCounted for Outer {
			fn irc_box(&self) -> &IrcBox<dyn IrcBoxed> {
				&self.inner.irc_box
			}
		}

		unsafe impl IntrusivelyCounted for Inner {
			fn irc_box(&self) -> &IrcBox<dyn IrcBoxed> {
				&self.irc_box
			}
		}

		// Create a composite structure with a member that includes an IrcBox.
		let mut outer = pin!(Lease::new(Outer::default()));

		// Derive a reference to the outmost instance.
		let outer_ref = NonNull::from(unsafe { outer.as_mut().project().get_unchecked_mut() });

		// Create an Irc-pointer to that composite structure.
		let irc1 = Irc::new(outer);

		// Project the Irc onto the inner one, dropping the outer reference.
		let irc2 = Irc::map(irc1, |inner| &inner.inner);
		let outer_ref = with_provenance(outer_ref, Irc::as_raw(&irc2));

		// Store it in the IrcBox.
		irc2.irc_box.outer_ref.set(Some(outer_ref));

		// Now drop the inner reference, leading to mut access of Outer.
		drop(irc2);
	}
}