odem_rs_core/simulator/

mod.rs

//! This module contains the heart of the library: the [Simulator] and the
//! [simulation context].
//!
//! Combined, these two structures implement *executor* and *spawner* for the
//! *future*-based implementation of coroutines. Conceptually, the executor is
//! responsible for selecting the next future to be executed, and the spawner
//! allows new futures to be injected into the executor.
//!
//! In our implementation, the executor is realized by the [Simulator], the
//! spawner is realized by the [simulation context], and futures are realized
//! through [Continuations](Continuation).
//!
//! [simulation context]: Sim

use core::{
	any::Any,
	cell::Cell,
	fmt,
	future::Future,
	num::NonZero,
	ops::Add,
	pin::{Pin, pin},
	ptr::NonNull,
	task::{Context, RawWaker, RawWakerVTable, Waker},
};

use crate::{
	Active, Puck,
	calendar::{Calendar, Partition},
	config::{Config, Time},
	continuation::{Continuation, Puck as ContPuck, Share, token},
	erased,
	error::{CausalityError, Deadlock},
	fsm::*,
	job::{Builder as JobBuilder, Job},
	ptr::{AsIrc, IntrusivelyCounted, Irc, IrcBox, IrcBoxed, Lease, LeasedMut},
};

pub mod details;
mod ops;

/// The type used to mark the agents of incoming continuations.
///
/// This is used to prevent interleaved execution of continuations from
/// different agents by enforcing the following property: all continuations
/// belonging to some agent `A` are sorted before any continuations belonging to
/// agent `B` iff any continuation of agent `A` is scheduled before all the
/// continuations of agent `B`.
pub(crate) type Mark = u64;

/// The precedence type for ordering the continuations that are executed as part
/// of the same agent.
///
/// The way this type works is by partitioning an integer range into a set of
/// non-overlapping intervals, each represented by the first integer that is
/// part of the respective interval. Initially, the interval contains the whole
/// range (`0..1`) but can be refined at every level into smaller subintervals
/// that are sorted according to their first integer and nesting level. This has
/// the effect of allowing subprocesses to inherit the precedence of their
/// parent and to recursively allow the same for their children. Since the
/// number of bits is finite, this is bound to fail at some level of nesting,
/// but it should not place undue limitations on the user.
pub type Prec = Partition<u32>;

/* ****************************************************** Simulation Function */

/// Simplified function that creates a new [default] configuration, runs the
/// provided `async` closure to completion, and returns it.
///
/// Use it in situations when both input and output of the simulation are
/// communicated through a configuration initialized to a default value.
///
/// [default]: Default
#[track_caller]
pub fn simulation<C, F>(main: F) -> Result<C, Deadlock<C::Time>>
where
	C: Config + Default,
	F: AsyncFnOnce(&Sim<C>),
{
	let sim = Simulator::default();
	sim.run(main).map(move |()| sim.into_inner())
}

/* **************************************************************** Simulator */

/// The principal driver of a simulation run which can be initialized using a
/// [configuration].
///
/// [configuration]: Config
#[derive(Copy, Clone, Default)]
pub struct Simulator<C>(C);

impl<C> Simulator<C> {
	/// Creates a new simulator from a configuration.
	pub const fn new(config: C) -> Self {
		Self(config)
	}

	/// Converts the simulator back into a configuration.
	pub fn into_inner(self) -> C {
		self.0
	}

	/// Grants shared access to the inner configuration.
	pub fn inner(&self) -> &C {
		&self.0
	}

	/// Grants mutable access to the inner configuration.
	pub fn inner_mut(&mut self) -> &mut C {
		&mut self.0
	}
}

impl<C: Config> Simulator<C> {
	/// Executes the simulation model provided in the closure and returns its
	/// result in the `Ok`-variant of a [Result].
	///
	/// The `Err`-variant is reserved for errors that occur during the
	/// event-loop. Only a [Deadlock] can be reported.
	#[track_caller]
	pub fn run<F, R>(&self, main: F) -> Result<R, Deadlock<C::Time>>
	where
		F: AsyncFnOnce(&Sim<C>) -> R,
	{
		// Create a fresh simulation context.
		// SAFETY: the reference to `self` outlives the simulation context.
		// (done to prevent mentioning the lifetime in the simulation context)
		let sim = pin!(unsafe { Sim::new(&self.0) });
		let sim = Irc::new(sim);

		// Register the simulation context thread-locally to support tracing.
		// The hook will take care of deregistration in its destructor.
		let _hook = erased::hook(sim.clone());

		// Bootstrap the root process.
		let share = pin!(Share::root(sim.clone()));
		let share = share.into_ref();

		#[cfg(feature = "tracing")]
		let span = tracing::error_span!(parent: None, "SimMain").entered();
		let root = pin!(JobBuilder::root().with_actions(main(&sim)).finish());
		let mut puck = Job::boot(root, share);

		#[cfg(feature = "tracing")]
		drop(span);

		// Schedule the root job at the current model-time.
		puck.as_irc().brand(|task, once| {
			let idle = task.token(once).into_idle().unwrap();
			sim.calendar().activate(task.clone(), idle);
		});

		{
			let waker = ShallowWaker::new(sim.clone());
			// SAFETY: the `ShallowWaker` is valid for the lifetime of the waker
			let waker = unsafe { waker.as_waker() };
			let mut cx = Context::from_waker(&waker);
			let root = puck.as_ref();

			// Run the simulator until the root job terminates.
			while let Some(cont) = sim.calendar.extract() {
				let root_terminated = cont.brand(|cont, once| {
					let next = cont.token(once).into_next().unwrap();
					let _span = cont.enter_span();
					let busy: token::Busy<'_> = cont.state().transition(next, ());

					// set the active continuation
					sim.active.set(Some(ContPuck::new(cont.clone(), &busy)));

					// poll the continuation and check whether the root continuation terminated
					cont.poll(busy, &mut cx).is_ready() && core::ptr::eq(cont.detach(), root)
				});

				if root_terminated {
					break;
				}
			}

			// Clear the active slot on exit.
			sim.active.set(None);
		}

		// Extract and return the result.
		puck.result().ok_or(Deadlock(sim.now()))
	}

	/// Executes the simulation model provided in the closure that has to
	/// return a [`Result`].
	///
	/// The result of a simulation run can also be a [`Deadlock`] if the
	/// scheduler detects that it is not possible to schedule any more
	/// continuations after some model-time has passed.
	#[track_caller]
	pub fn with_result<F, T, E>(&self, main: F) -> Result<T, E>
	where
		F: AsyncFnOnce(&Sim<C>) -> Result<T, E>,
		E: From<Deadlock<C::Time>>,
	{
		self.run(main).unwrap_or_else(|err| Err(E::from(err)))
	}
}

impl<C> From<C> for Simulator<C> {
	fn from(config: C) -> Self {
		Self::new(config)
	}
}

/* ******************************************************* Simulation Context */

/// Simulation context that can be used to advance the [model-time], spawn
/// additional [Processes] or [Jobs], access [global data], the [current]
/// model-time and the [active] process, among other things.
///
/// It can only be accessed during a simulation run.
///
/// [Processes]: crate::agent::Agent
/// [Jobs]: Job
/// [global data]: Config::Data
/// [model-time]: Config::Time
/// [current]: Sim::now
/// [active]: Sim::active
pub struct Sim<C: ?Sized + Config = ()> {
	/// The event calendar.
	calendar: Calendar<C>,
	/// The currently active continuation.
	active: Cell<Option<ContPuck<C>>>,
	/// The inner box to enable intrusive reference counting.
	sim_box: IrcBox,
	/// A global counter for unique agent-IDs.
	#[cfg(not(feature = "alloc"))]
	pid_gen: Cell<NonZero<usize>>,
	/// A per-agent-type counter for unique agent-IDs.
	#[cfg(feature = "alloc")]
	pid_gen: core::cell::RefCell<hashbrown::HashMap<core::any::TypeId, NonZero<usize>>>,
	/// Pointer to the user-supplied configuration.
	config: NonNull<C>,
}

impl<C: Config> Sim<C> {
	/// Creates a new simulation context from a [configuration](Config).
	///
	/// # Safety
	/// The caller has to ensure that the reference to the configuration
	/// outlives the created simulation context.
	///
	/// This could also be solved via the borrow checker but would require the
	/// introduction of an additional lifetime which makes referring to the
	/// simulation context in function signatures quite a bit less ergonomic.
	#[track_caller]
	unsafe fn new<'p>(config: &C) -> Lease<'p, Self> {
		Lease::new(Sim {
			calendar: Calendar::new(config),
			active: Cell::new(None),
			sim_box: IrcBox::default(),
			pid_gen: {
				#[cfg(not(feature = "alloc"))]
				{
					Cell::new(NonZero::<usize>::MIN)
				}
				#[cfg(feature = "alloc")]
				{
					core::cell::RefCell::new(hashbrown::HashMap::new())
				}
			},
			config: NonNull::from(config),
		})
	}

	/// Gets a [simulation context] from a waker.
	///
	/// This can be used to recover the current simulation context from the
	/// context argument if the configuration is known, making it possible to
	/// distribute the simulation context by piggybacking on Rust's implicit
	/// [Context].
	///
	/// [simulation context]: Sim
	pub fn from_context(context: &Context<'_>) -> Option<Irc<Self>> {
		let waker = context.waker();

		// comparing for pointer equality is enough because there is only
		// one vtable for all simulators
		if core::ptr::eq(waker.vtable(), &SIM_VTABLE) {
			// SAFETY: it's a shallow waker, so it is safe to reconstruct the
			// simulation context
			let waker = unsafe { &*(waker.data() as *const ShallowWaker) };
			waker.0.clone().downcast::<Self>().ok()
		} else {
			// none of our wakers
			None
		}
	}
}

impl<C: ?Sized + Config> Sim<C> {
	/// Removes the currently active continuation.
	pub(crate) fn unslot<'brand>(
		&self,
		task: &Continuation<'brand, C>,
		busy: token::Busy<'brand>,
	) -> token::Idle<'brand> {
		let _puck = self.active.take();

		// sanity check in debug mode
		debug_assert!(
			_puck.map(|puck| puck.is_same(task)).unwrap_or(true),
			"unslotted task in state `Busy` was not in the active slot"
		);

		task.state().transition(busy, ())
	}

	/// Returns a shared reference to the configuration.
	pub fn config(&self) -> &C {
		// SAFETY: The configuration outlives the simulation context per the
		// invariant on the unsafe `Self::new`.
		unsafe { self.config.as_ref() }
	}

	/// Returns a reference to the shared data.
	pub fn global(&self) -> &C::Data {
		self.config().global_data()
	}

	/// Returns the [Puck] of the currently active continuation.
	pub fn active(&self) -> ContPuck<C> {
		let active = self.active.take();
		let result = active.clone();
		self.active.set(active);
		result.expect("no active continuation")
	}

	/// Changes the rank of the currently active [`Agent`](crate::agent::Agent).
	///
	/// The new rank takes immediate effect and causes the rearrangement
	/// of all jobs currently scheduled for the active agent, both in the
	/// present and future.
	///
	/// Lowering the rank of the active `Agent` can lead to another `Agent`
	/// gaining control if one with a (now) higher rank has jobs scheduled at
	/// the current model time. This change takes effect once the currently
	/// active agent suspends.
	pub fn update_rank(&self, rank: C::Rank) {
		self.active().share().update_rank(rank);
	}

	/// Returns the current model time.
	pub fn now(&self) -> C::Time {
		self.calendar().now()
	}

	/// Schedules the currently active continuation at a later model time.
	///
	/// Panics if no continuation is currently active or if the proposed model time has
	/// already passed. For a non-panicking version of this method, see
	/// [`try_advance`].
	///
	/// [`try_advance`]: Self::try_advance
	pub fn advance(&self, dt: impl Into<C::Time>) -> impl Future<Output = ()> + '_
	where
		C::Time: Add<Output = C::Time>,
	{
		ops::TryAdvance::new(self, dt.into()).unwrap()
	}

	/// Schedules the currently active continuation at a later model time.
	///
	/// This method has the same function as [`advance`] but takes an absolute
	/// time-point for the proposed reactivation time rather than a relative
	/// one.
	///
	/// For a non-panicking alternative, see [`try_advance_to`].
	///
	/// [`advance`]: Self::advance
	/// [`try_advance_to`]: Self::try_advance_to
	pub fn advance_to(&self, time: impl Into<C::Time>) -> impl Future<Output = ()> + '_ {
		ops::TryAdvanceTo::new(self, time.into()).unwrap()
	}

	/// Attempts to advance the currently active continuation to a later model time.
	///
	/// Fails with an `Err`-result if either no continuation is active or if the
	/// proposed model time has already passed.
	pub fn try_advance(
		&self,
		dt: impl Into<C::Time>,
	) -> impl Future<Output = Result<(), CausalityError<C::Time>>> + '_
	where
		C::Time: Add<Output = C::Time>,
	{
		ops::TryAdvance::new(self, dt.into())
	}

	/// Attempts to advance the currently active continuation to a later model time.
	///
	/// This method has the same function as [`try_advance`] but takes an
	/// absolute time-point for the proposed reactivation time rather than a
	/// relative one.
	///
	/// [`try_advance`]: Self::try_advance
	pub fn try_advance_to(
		&self,
		time: impl Into<C::Time>,
	) -> impl Future<Output = Result<(), CausalityError<C::Time>>> + '_ {
		ops::TryAdvanceTo::new(self, time.into())
	}

	/// Activates a new continuation by binding it to the simulation context,
	/// scheduling it at the current model time, and returning a [`Puck`] to it.
	///
	/// # Example Usage
	///
	/// You will usually want to [`pin`] the active object to pass it into the
	/// function like this:
	///
	/// ```
	/// # use {core::pin::pin, odem_rs_core::{simulator::Sim, job::Job}};
	/// # async fn sim_main(sim: &Sim) {
	/// let job = pin!(Job::new(async { /* do important work */ }));
	/// let puck = sim.activate(job);
	/// # }
	/// ```
	///
	/// # Note
	///
	/// The function expects a pinned mutable reference to a [`Lease`] to an
	/// active object to utilize the borrow-checker to prevent certain
	/// usage errors statically. Without going into too much technical detail,
	/// the argument is made into the equivalent of a reference with move
	/// semantics. Once this function has been called, it is *as if* the
	/// value has been moved into the function, preventing access to it from the
	/// caller's context, even though only a reference has been moved.
	///
	/// Specifically, the following error is prevented:
	/// ```compile_fail,E0505
	/// # use {core::pin::pin, odem_rs_core::{simulator::Sim, job::Job}};
	/// # async fn sim_main(sim: &Sim) {
	/// let mut job = pin!(Job::new(async { /* do important work */ }));
	/// sim.activate(job.as_mut()); // allowed
	/// // later
	/// sim.activate(job);          // error: cannot move out of `job` because it is borrowed
	/// # }
	/// ```
	///
	/// Please refer to the documentation of `Lease` for technical details.
	pub fn activate<'a, A>(&'a self, actions: Pin<LeasedMut<'a, A>>) -> A::Puck<'a>
	where
		A: Active<C>,
	{
		// bind the simulation context
		let puck = self.bind(actions);

		// activate the puck
		puck.as_irc().brand(|task, once| {
			// successfully bound continuations are in state 'Idle'
			let idle = task.token(once).into_idle().unwrap();

			// set the correct span
			let _span = task.enter_span();

			// schedule the puck
			self.calendar().activate(task.clone(), idle);

			puck
		})
	}

	/// Activates a new continuation by binding it to the simulation context,
	/// scheduling it at a future model time, and returning a [`Puck`] to it.
	///
	/// The function panics if the model time for activation has already passed.
	/// For a non-panicking version of this method, see [`try_schedule`].
	///
	/// # Example Usage
	///
	/// You will usually want to [`pin`] the active object to pass it into the
	/// function like this:
	///
	/// ```
	/// # use {core::pin::pin, odem_rs_core::{simulator::Sim, job::Job}};
	/// # async fn sim_main(sim: &Sim) {
	/// let job = pin!(Job::new(async { /* do important work */ }));
	/// let puck = sim.schedule(job, 2.0); // schedule in 2 units of time
	/// # }
	/// ```
	///
	/// # Note
	///
	/// The notes from [`Self::activate`] apply here as well.
	///
	/// [`try_schedule`]: Self::try_schedule
	#[track_caller]
	pub fn schedule<'a, A>(
		&'a self,
		actions: Pin<LeasedMut<'a, A>>,
		dt: impl Into<C::Time>,
	) -> A::Puck<'a>
	where
		A: Active<C>,
		C::Time: Add<Output = C::Time>,
	{
		self.try_schedule(actions, dt).unwrap_or_else(|err| {
			#[cfg(feature = "tracing")]
			tracing::error!(%err);
			panic!("{}", err);
		})
	}

	/// Activates a new continuation, schedule it at a later model time, and
	/// returns a [`Puck`] to it.
	///
	/// Doesn't panic if the proposed model time has already passed, returning
	/// a [`CausalityError`] instead. For the panicking version, see
	/// [`schedule`].
	///
	/// [`schedule`]: Self::schedule
	pub fn try_schedule<'a, A>(
		&'a self,
		actions: Pin<LeasedMut<'a, A>>,
		dt: impl Into<C::Time>,
	) -> Result<A::Puck<'a>, CausalityError<C::Time>>
	where
		A: Active<C>,
		C::Time: Add<Output = C::Time>,
	{
		use core::cmp::Ordering::*;
		let time: C::Time = self.now() + dt.into();

		match time.partial_cmp(&self.now()) {
			None | Some(Less) => Err(CausalityError {
				cause: self.now(),
				effect: time,
			}),
			Some(order) => {
				// bind the simulation context
				let puck = self.bind(actions);

				// schedule or activate the puck
				puck.as_irc().brand(|task, once| {
					// successfully bound continuations are in state 'Idle'
					let idle = task.token(once).into_idle().unwrap();

					// set the correct span
					let _span = task.enter_span();

					// schedule the puck
					let calendar = self.calendar();

					let irc = task.clone();
					if order == Greater {
						calendar.schedule(irc, idle, time);
					} else {
						calendar.activate(irc, idle);
					}

					Ok(puck)
				})
			}
		}
	}

	/// Returns an awaitable [`Future`] that will suspend the currently executed
	/// continuation and reschedule it after every other continuation scheduled
	/// at the current model time had a chance to run.
	///
	/// The free function [`defer`] has the same semantics but reconstructs the
	/// simulation context through dynamic dispatching.
	///
	/// [`defer`]: crate::ops::defer
	pub fn defer(&self) -> impl Future<Output = ()> + '_ {
		ops::Defer::new(self)
	}

	/// Generates a new, simulation-run-unique agent id.
	pub(crate) fn pid_gen<I: 'static>(&self) -> NonZero<usize> {
		#[cfg(not(feature = "alloc"))]
		// increment the global agent counter to get a new pid
		{
			let next = self.pid_gen.get();
			self.pid_gen
				.set(next.checked_add(1).unwrap_or(NonZero::<usize>::MIN));
			next
		}

		#[cfg(feature = "alloc")]
		// look up the agent-type-specific pid to increment
		{
			let mut pid_gen = self.pid_gen.borrow_mut();
			let counter = pid_gen
				.entry(core::any::TypeId::of::<I>())
				.or_insert(NonZero::<usize>::MIN);
			let next = *counter;
			*counter = next.checked_add(1).unwrap_or(NonZero::<usize>::MIN);
			next
		}
	}

	/// Returns a reference to the [event calendar].
	///
	/// [event calendar]: crate::DefaultPlan
	pub(crate) fn calendar(&self) -> &Calendar<C> {
		&self.calendar
	}

	/// Activates the active object, returning a [`Puck`] on success and an
	/// [`Error`] if the operation failed.
	///
	/// This is slightly more complicated than simply calling [`bind`] on the
	/// [`Active`]-trait because we also need to store the *actual* type of the
	/// instance being activated to allow dynamic dispatching to it later.
	///
	/// [`bind`]: Active::bind
	fn bind<'a, A>(&'a self, mut actions: Pin<LeasedMut<'a, A>>) -> A::Puck<'a>
	where
		A: Active<C>,
	{
		// create a raw reference to the pinned continuation;
		// this will be the most specialized version, enabling dynamic dispatch
		let vtab = NonNull::from(unsafe { actions.as_mut().project().get_unchecked_mut() });

		// create a reference to the context of the currently active object
		// TODO: this seems unsafe due to the lifetime extension
		let active = unsafe { &*self.active.as_ptr() }.as_ref().unwrap();
		let share = active.share();

		// create the puck by binding
		let puck = Active::bind(actions, share);

		// override the vptr
		unsafe {
			puck.as_ref().set_vptr(vtab);
		}

		// return the puck
		puck
	}
}

// Configuration-erased dyn-safe version of the simulation context.
impl<C: Config> erased::Sim for Sim<C> {
	fn active(&self) -> Option<erased::ContPuck> {
		let active = self.active.take();
		self.active.set(active.clone());
		active.map(Into::into)
	}

	fn now(&self) -> &dyn fmt::Display {
		// SAFETY: repr(transparent) guarantees that the two structures are laid
		// out identically in memory, making this transmutation safe.
		unsafe { core::mem::transmute::<&Sim<C>, &DisplayHelper<C>>(self) }
	}

	fn global(&self) -> &dyn Any {
		self.global()
	}

	fn waker(&self) -> RawWaker {
		self.active().into_waker()
	}

	fn defer(&self) {
		if let Some(active) = self.active.take() {
			active.into_inner().brand(|active, once| {
				let busy = active
					.token(once)
					.into_busy()
					.expect("active continuation should be 'Busy'");
				self.calendar().defer(active.clone(), busy);
			});
		}
	}
}

impl<C: ?Sized + Config> fmt::Debug for Sim<C>
where
	C::Plan: fmt::Debug,
{
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let mut debug = f.debug_struct("Sim");
		let active = self.active.take();
		self.active.set(active.clone());

		debug
			.field("active", &active)
			.field("calendar", &self.calendar)
			.finish()
	}
}

/// *Newtype* adapter that implements [Display](fmt::Display) for a simulation
/// context by displaying the current model-time.
#[repr(transparent)]
struct DisplayHelper<C: Config>(Sim<C>);

impl<C: Config> fmt::Display for DisplayHelper<C> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		self.0.now().format(f)
	}
}

unsafe impl<C: ?Sized + Config> IntrusivelyCounted for Sim<C> {
	#[inline(always)]
	fn irc_box(&self) -> &IrcBox<dyn IrcBoxed> {
		&self.sim_box
	}
}

/* ************************************************************ Shallow Waker */

/// Internal waker that can only be used to create a copy of the currently
/// active continuation or to defer it.
struct ShallowWaker(Irc<dyn erased::Sim>);

/// Virtual function table for the [simulation executor](Simulator), independent
/// of the [configuration](Config).
///
/// This construction ensures that there is a globally unique address for the
/// simulators, allowing us to restore the [simulation context](Sim) from the
/// context and the configuration.
static SIM_VTABLE: RawWakerVTable = RawWakerVTable::new(
	ShallowWaker::clone,
	ShallowWaker::nop,
	ShallowWaker::wake_by_ref,
	ShallowWaker::nop,
);

impl ShallowWaker {
	/// Creates a [Waker] for a simulator that can be cloned into [Puck]
	/// wakers. Waking this waker has the effect of deferring the currently
	/// active continuation, i.e., the continuation is reawakened at the current
	/// model time after all the other continuations have run.
	fn new<C: Config>(sim: Irc<Sim<C>>) -> Self {
		// erase the concrete configuration
		Self(Irc::map(sim, |inner| inner as &dyn erased::Sim))
	}

	/// Constructs a [Waker] from the reference of a shallow waker.
	///
	/// # Safety
	/// The caller is responsible to ensure that the reference to the shallow
	/// waker is valid for the lifetime of the waker.
	unsafe fn as_waker(&self) -> Waker {
		unsafe {
			let waker = RawWaker::new(self as *const Self as *const (), &SIM_VTABLE);
			Waker::from_raw(waker)
		}
	}

	/// No-op waker-function.
	fn nop(_: *const ()) {}

	/// Constructs a [RawWaker] for the active process.
	unsafe fn clone(this: *const ()) -> RawWaker {
		unsafe {
			// perform a dynamic dispatch to recover the configuration
			(*(this as *const Self)).0.waker()
		}
	}

	/// Defers the currently active process.
	unsafe fn wake_by_ref(this: *const ()) {
		unsafe {
			// perform a dynamic dispatch to recover the configuration
			(*(this as *const Self)).0.defer();
		}
	}
}