odem_rs_util/

pool.rs

//! This module defines an arena allocator for [`Active`] simulation objects
//! that performs automatic garbage collection on terminated and unreferenced
//! instances.
//!
//! # Overview
//!
//! The primary components are:
//!
//! - [`Pool`]: The main arena allocator. It manages a backing store and a
//!   list of recyclable memory slots. It must be pinned to guarantee stable
//!   addresses for the active entities it allocates.
//! - [`Storage`]: A trait that abstracts over the memory source for the pool,
//!   allowing for both fixed-size (`no_alloc`) and dynamically growing
//!   (`alloc`) backing stores.
//! - [`Gc<A>`]: A wrapper around an `Active` type `A`. When the last
//!   reference to a `Gc` is dropped, it automatically returns its memory slot
//!   to the pool for recycling instead of deallocating it.
//!
//! # Examples
//!
//! Creating a statically sized pool for `no_alloc` environments.
//! ```
//! # use core::pin::pin;
//! # use odem_rs_util::pool::Pool;
//! # use odem_rs_core::{config::Config, job::Job, simulator::Sim};
//! # #[derive(Default, Config)] struct MyConfig;
//! async fn sim_main(sim: &Sim<MyConfig>) {
//!     // Create a pool for up to five elements.
//!     let pool = pin!(Pool::fixed::<5>());
//!     
//!     // Allocate a new Job and activate it in the simulation.
//!     // The Gc wrapper ensures its memory is recycled when it's done.
//!     sim.activate(pool.alloc(Job::new(async {})));
//! }
//! ```
//!
//! Creating a dynamically growing pool for `alloc` environments.
//! ```
//! # use core::pin::pin;
//! # use odem_rs_util::pool::Pool;
//! # use odem_rs_core::{config::Config, job::Job, simulator::Sim};
//! # #[derive(Default, Config)] struct MyConfig;
//! async fn sim_main(sim: &Sim<MyConfig>) {
//!     // Create a dynamically growing pool for jobs.
//!     let pool = pin!(Pool::dynamic().with(Job::new));
//!     
//!     // Allocate a new Job and activate it in the simulation.
//!     // The Gc wrapper ensures its memory is recycled when it's done.
//!     sim.activate(pool.alloc(async {}));
//! }
//! ```
//!
//! [`Arena`]: typed_arena::Arena
#![allow(missing_docs)]

use core::{
	cell::Cell,
	pin::Pin,
	ptr::NonNull,
	task::{Context, Poll},
};

use intrusive_collections::{SinglyLinkedList, SinglyLinkedListLink, UnsafeRef};

use odem_rs_core::{Active, Dispatch, ExitStatus, config::Config, continuation::Share, ptr::Lease};

use crate::error::InsufficientSpace;

#[doc(hidden)]
mod adapter;

mod backing;

#[cfg(test)]
mod tests;

pub use adapter::*;
pub use backing::*;

/* ********************************************************** Arena Allocator */

/// An arena allocator for [Active] simulation objects that performs automatic
/// garbage collection on terminated and unreferenced instances.
///
/// This pool manages a [`Storage`] backend and a `reclaim` list of memory
/// slots that are ready to be reused.
///
/// # Safety Notes
///
/// Rearranging the fields here can lead to undefined behavior upon dropping of
/// the pool. Specifically, dropping `storage` before `reclaim` will invalidate
/// the slots contained in `reclaim`, leading to undefined behavior when
/// `reclaim` is dropped later.
#[pin_project::pin_project]
pub struct Pool<A, S: ?Sized = dyn Storage<Slot = Gc<A>>, F = ()> {
	/// A linked list of terminated instances ready to be recycled.
	reclaim: Cell<SinglyLinkedList<GcAdapter<A>>>,
	/// A constructor function that constructs the value before it is placed
	/// inside a slot of the storage.
	cons: F,
	/// A backing store used to allocate more instances.
	#[pin]
	storage: S,
}

impl<A> Pool<A> {
	/// Creates a new pool given a backing store.
	pub fn new<S>(storage: S) -> Pool<A, S> {
		Pool {
			reclaim: Cell::new(SinglyLinkedList::new(GcAdapter::new())),
			cons: (),
			storage,
		}
	}

	/// Creates a pool with a fixed-size backing store.
	///
	/// This pool is suitable for `no_std` + `no_alloc` environments. It will
	/// panic if an allocation is attempted when the pool is full.
	pub fn fixed<const N: usize>() -> Pool<A, Array<Gc<A>, N>> {
		Pool::new(Array::new())
	}

	/// Creates a pool with a dynamically growing backing store.
	#[cfg(feature = "alloc")]
	#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
	pub fn dynamic() -> Pool<A, typed_arena::Arena<Gc<A>>> {
		Pool::new(typed_arena::Arena::new())
	}

	/// Creates a dynamically growing pool with a specified initial capacity.
	#[cfg(feature = "alloc")]
	#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
	pub fn with_capacity(n: usize) -> Pool<A, typed_arena::Arena<Gc<A>>> {
		Pool::new(typed_arena::Arena::with_capacity(n))
	}
}

impl<A, S: Storage<Slot = Gc<A>>> Pool<A, S> {
	/// Sets a new constructor function for the pool.
	pub fn with<'p, B, F>(self, cons: F) -> Pool<A, S, F>
	where
		F: Fn(B) -> Lease<'p, A>,
		A: 'p,
	{
		Pool {
			reclaim: self.reclaim,
			cons,
			storage: self.storage,
		}
	}
}

impl<A, S: ?Sized + Storage<Slot = Gc<A>>, F> Pool<A, S, F> {
	/// Attempts to allocate a new, [`Active`] object.
	///
	/// This method first tries to recycle a slot from the reclaim list. If the
	/// list is empty, it requests a new slot from the backing store.
	///
	/// This function is panic-safe. If the constructor function `cons` panics,
	/// the slot (whether recycled or newly allocated) is safely returned to
	/// the reclaim list, preventing memory leaks.
	///
	/// Returns a `Result` containing the pinned, allocated object, or the
	/// original `value` if allocation fails (e.g., the backing store is full).
	#[allow(clippy::mut_from_ref)]
	pub fn try_alloc<'a, 'p, B>(
		self: &'a Pin<&mut Self>,
		value: B,
	) -> Result<Pin<&'a mut Lease<'a, Gc<A>>>, InsufficientSpace<B>>
	where
		F: sealed::Cons<B, Lease<'p, A>>,
		A: 'p,
	{
		use scopeguard::{ScopeGuard, guard};

		// Try to get a slot from either the reclaim list or the backing store.
		let mut list = self.reclaim.take();
		let slot = list.pop_front();
		self.reclaim.set(list);

		let slot = match slot {
			Some(item) => {
				// SAFETY: The `UnsafeRef` points to a valid `Gc` in storage.
				unsafe { Pin::new_unchecked(&mut *UnsafeRef::into_raw(item)) }
			}
			None => {
				// The reclaim list was empty, so try the backing store.
				match self.as_ref().project_ref().storage.try_alloc(Gc::new()) {
					Some(new_slot) => new_slot,
					None => return Err(InsufficientSpace(value)),
				}
			}
		};

		// We have a slot. Guard it to ensure it's reclaimed on panic.
		let mut guard = guard(slot, |slot| {
			// This closure runs if a panic occurs during construction.
			let mut list = self.reclaim.take();
			// SAFETY: `slot` points to a valid memory location that was just
			// allocated but is being abandoned. It's safe to add it to the
			// reclaim list for future use.
			list.push_front(unsafe { UnsafeRef::from_raw(slot.get_unchecked_mut() as *mut _) });
			self.reclaim.set(list);
		});

		// Perform the potentially panicking operation.
		guard
			.as_mut()
			.fill(&self.reclaim, self.cons.cons(value).into_inner());

		// Success! Disarm the guard and get the slot back.
		let slot = ScopeGuard::into_inner(guard);

		// SAFETY: Adding the `Lease` here is safe, because this is the
		// only pinned reference to this slot in existence.
		Ok(unsafe { Lease::unchecked_new(slot) })
	}

	/// Allocates a new, [`Active`] object, panicking if allocation fails.
	///
	/// This is a convenience wrapper around [`try_alloc`](Self::try_alloc) that
	/// panics on allocation failure. This is often acceptable in simulations
	/// where the pool size is expected to be sufficient.
	#[allow(clippy::mut_from_ref)]
	pub fn alloc<'a, 'p, B>(self: &'a Pin<&mut Self>, value: B) -> Pin<&'a mut Lease<'a, Gc<A>>>
	where
		F: sealed::Cons<B, Lease<'p, A>>,
		A: 'p,
	{
		self.try_alloc(value).unwrap()
	}

	/// Returns the total number of allocated slots.
	pub fn len(&self) -> usize {
		let list = self.reclaim.take();
		let reclaimed = list.iter().count();
		self.reclaim.set(list);

		self.storage.len() - reclaimed
	}

	/// Returns `true` if no slots have been allocated.
	pub fn is_empty(&self) -> bool {
		self.len() == 0
	}
}

impl<A, S: Storage<Slot = Gc<A>>> From<S> for Pool<A, S> {
	fn from(storage: S) -> Self {
		Pool::new(storage)
	}
}

impl<A, S: Storage<Slot = Gc<A>> + Default> Default for Pool<A, S> {
	fn default() -> Self {
		Pool::new(S::default())
	}
}

/* ****************************************************** `Cons` Helper-Trait */

#[doc(hidden)]
mod sealed {
	pub trait Cons<S, T> {
		fn cons(&self, value: S) -> T;
	}

	// Implement the identity function for the empty tuple.
	// This is done to prevent wasting memory for the standard case without
	// a dedicated constructor function.
	impl<T> Cons<T, T> for () {
		#[inline]
		fn cons(&self, value: T) -> T {
			value
		}
	}

	// Blanket implement the trait for `Fn` closures.
	impl<S, T, F> Cons<S, T> for F
	where
		F: Fn(S) -> T,
	{
		#[inline]
		fn cons(&self, value: S) -> T {
			self(value)
		}
	}
}

/* ************************************************************ Storage Trait */

/// A trait for abstracting over the memory backend of a [`Pool`].
///
/// This allows the `Pool` to be generic over its memory source, supporting
/// both fixed-size arrays for `no_alloc` environments and dynamic arenas
/// where a heap is available.
pub trait Storage {
	/// Slot-type of the memory pool.
	///
	/// Should be [`Gc<A>`] for some [`Active`] `A`.
	type Slot;

	/// Allocates a new object, panicking if the backing store is full.
	#[allow(clippy::mut_from_ref)]
	fn alloc(self: Pin<&Self>, value: Self::Slot) -> Pin<&mut Self::Slot> {
		self.try_alloc(value).expect("insufficient storage space")
	}

	/// Attempts to allocate a new object, returning an error if the store is
	/// full.
	#[allow(clippy::mut_from_ref)]
	fn try_alloc(self: Pin<&Self>, value: Self::Slot) -> Option<Pin<&mut Self::Slot>>;

	/// Returns the number of items currently allocated in the store.
	fn len(&self) -> usize;

	/// Returns `true` if no items have been allocated from the store.
	fn is_empty(&self) -> bool {
		self.len() == 0
	}
}

/* ************************************************** Garbage Collecting Type */

/// A wrapper that enables an [`Active`] type to be managed by a [`Pool`].
///
/// This enum has two states:
/// - `Active`: The wrapper contains a live `Active` entity.
/// - `Reclaim`: The entity has been dropped, and this `Gc` now acts as a node
///   in the pool's intrusive reclaim list.
#[pin_project::pin_project(
	project = GcProj,
	project_ref = GcProjRef,
	project_replace = GcProjReplace
)]
pub enum Gc<A> {
	/// The state of a `Gc` when it contains a live, active entity.
	Active {
		/// A pointer to the pool's reclaim list.
		reclaim: NonNull<Cell<SinglyLinkedList<GcAdapter<A>>>>,
		/// The pinned active entity itself.
		#[pin]
		active: A,
	},
	/// The state of a `Gc` after its entity has been reclaimed. It now holds
	/// only a link for the intrusive list.
	Reclaim(SinglyLinkedListLink),
}

impl<A> Gc<A> {
	/// Creates a slot in the empty `Reclaim` state.
	fn new() -> Self {
		Self::Reclaim(SinglyLinkedListLink::new())
	}

	/// Creates a new garbage-collected instance in the `Active` state.
	fn fill<'a>(
		mut self: Pin<&'a mut Self>,
		reclaim: &'a Cell<SinglyLinkedList<GcAdapter<A>>>,
		active: A,
	) {
		self.set(Self::Active {
			reclaim: NonNull::from(reclaim),
			active,
		});
	}
}

impl<A: Dispatch> Dispatch for Gc<A> {
	/// Polls the underlying `Active` entity.
	fn poll(self: Pin<&Self>, cx: &mut Context<'_>) -> Poll<ExitStatus> {
		match self.project_ref() {
			GcProjRef::Active { active, .. } => active.poll(cx),
			GcProjRef::Reclaim { .. } => {
				// This state should be impossible to reach during a poll, as an
				// entity is only reclaimed after it has finished execution.
				unreachable!("Cannot poll a reclaimed Gc instance");
			}
		}
	}

	/// Reclaims the entity, returning its memory to the pool.
	///
	/// This method is called when the last `Irc` to this `Gc` is dropped.
	/// It transitions the enum from the `Active` state to the `Reclaim` state.
	/// This involves dropping the contained `Active` entity and then pushing
	/// the `Gc` wrapper itself onto the pool's intrusive reclaim list.
	fn reclaim(mut self: Pin<&mut Self>) {
		// Atomically replace the `Active` variant with `Reclaim`. The
		// `project_replace` gives us ownership of the fields from the previous
		// `Active` state. The `active` field will be dropped here, running its
		// destructor as intended.
		let prev = self.as_mut().project_replace(Gc::new());

		match prev {
			GcProjReplace::Active { reclaim, .. } => {
				// Now that the entity is dropped, add its memory slot (which is
				// `self`, now in the `Reclaim` state) to the reclaim list.
				let reclaim = unsafe { reclaim.as_ref() };

				// Take temporary ownership of the reclaim list.
				let mut list = reclaim.take();

				// SAFETY: We have mutable access to `self`, and the adapter
				// correctly calculates the offset to the link field.
				list.push_front(unsafe { UnsafeRef::from_raw(self.get_unchecked_mut() as *mut _) });

				// Return the modified list to the pool.
				reclaim.set(list);
			}
			GcProjReplace::Reclaim { .. } => {
				unreachable!("Cannot reclaim a Gc instance that has already been reclaimed");
			}
		}
	}
}

impl<C: ?Sized + Config, A: Active<C>> Active<C> for Gc<A> {
	type Output = A::Output;
	type Puck<'p>
		= A::Puck<'p>
	where
		A: 'p;

	/// Binds the underlying active entity to the simulation context.
	fn bind<'p>(this: Pin<&'p mut Lease<'p, Self>>, sctx: &'p Share<C>) -> Self::Puck<'p> {
		// We project through the `Gc` wrapper to get to the inner `active` field.
		Active::bind(
			this.map_pin_mut(|inner| match inner.project() {
				GcProj::Active { active, .. } => active,
				GcProj::Reclaim { .. } => {
					unreachable!("Cannot bind a reclaimed Gc instance");
				}
			}),
			sctx,
		)
	}
}