odem_rs_sync/

fork.rs

//! Provides support for agent-internal parallelism in simulation contexts.
//!
//! This module implements a forking mechanism inspired by SLX, enabling
//! concurrent execution of futures.
//!
//! # Features
//! - **Join Futures**: Use [`and`] to wait for all futures to complete.
//! - **Select Futures**: Use [`or`] to resolve when any one future completes.
//! - Supports nesting of combinators for complex workflows.
//!
//! [`and`]: Fork::and
//! [`or`]: Fork::or

pub use {
	all_of::{AllOf, Joinable},
	any_of::{AnyOf, Pickable},
};

use all_of::NotifyQuorum;
use any_of::NotifyOnExit;
use core::{future::IntoFuture, panic::Location, pin::Pin};
use odem_rs_core::{Active, config::Config, job::Job, ptr::Lease, simulator::Sim};

mod all_of;
mod any_of;

/* ***************************************************** Fork Extension Trait */

/// Extension trait for [simulation contexts] to support the [fork]-operation.
///
/// [simulation contexts]: Sim
/// [fork]: ForkExt::fork
pub trait ForkExt {
	/// Associated [simulation configuration] type.
	///
	/// [simulation configuration]: Config
	type Config: ?Sized + Config;

	/// Creates a future combinator that enables agent-internal parallelism
	/// by combining futures for concurrent execution within the simulation
	/// context.
	///
	/// The `fork` method allows developers to do either:
	///
	/// 1. **Wait for any one of the futures to complete** ([`or`] combinator).
	/// 2. **Wait for all futures to complete** ([`and`] combinator).
	///
	/// ## Usage Notes
	/// - The `or` combinator requires all futures to share the same return
	///   type. If multiple futures could complete at the same simulation time,
	///   precedence is given to those chained earlier.
	/// - The `and` combinator collects the results of all futures into a tuple,
	///   preserving their order in the chain.
	/// - Nested combinators can be used to create advanced logic.
	///
	/// ## Examples
	///
	/// ### Using `or` to Wait for the First Completion
	/// ```
	/// # use {odem_rs_core::simulator::Sim, odem_rs_sync::fork::ForkExt as _};
	/// # async fn sim_main(sim: &Sim) {
	/// let result = sim
	///     .fork(async { 1 })
	///     .or(async { 2 })
	///     .or(async { 3 })
	///     .await;
	///
	/// assert_eq!(result, 1);
	/// # }
	/// ```
	///
	/// ### Using `and` to Wait for All Completions
	/// ```
	/// # use {odem_rs_core::simulator::Sim, odem_rs_sync::fork::ForkExt as _};
	/// # async fn sim_main(sim: &Sim) {
	/// let results = sim
	///     .fork(async { true })
	///     .and(async { 2 })
	///     .and(async { "3" })
	///     .await;
	///
	/// assert_eq!(results, (true, 2, "3"));
	/// # }
	/// ```
	///
	/// ### Combining Nested Futures
	/// ```
	/// # use {odem_rs_core::simulator::Sim, odem_rs_sync::fork::ForkExt as _};
	/// # async fn sim_main(sim: &Sim) {
	/// let result = sim
	///     .fork(async { (1, "One") })
	///     .or(sim.fork(async { 2 }).and(async { "Two" }))
	///     .await;
	///
	/// assert_eq!(result, (1, "One"));
	/// # }
	/// ```
	///
	/// [`or`]: Fork::or
	/// [`and`]: Fork::and
	fn fork<F: IntoFuture>(&self, future: F) -> Fork<'_, Self::Config, F>;
}

// implement the trait for all valid simulation contexts
impl<C: ?Sized + Config> ForkExt for Sim<C> {
	type Config = C;

	#[track_caller]
	fn fork<F: IntoFuture>(&self, fut: F) -> Fork<'_, C, F> {
		Fork {
			sim: self,
			fut,
			loc: Location::caller(),
		}
	}
}

/* ********************************************************************* Fork */

/// A builder-type structure that allows assembling arbitrary many futures into
/// a concurrently executed choice or collection.
pub struct Fork<'s, C: ?Sized + Config, F> {
	/// Associated simulation context.
	sim: &'s Sim<C>,
	/// A future to await.
	fut: F,
	/// A reference to the source location of the future.
	loc: &'static Location<'static>,
}

impl<'s, C: ?Sized + Config, F: IntoFuture + 's> Fork<'s, C, F> {
	/// Converts the `Fork` into a [`AnyOf`]-type of choice.
	#[track_caller]
	pub fn or<G: IntoFuture<Output = F::Output> + 's>(
		self,
		fut: G,
	) -> AnyOf<'s, C, impl Pickable<C, Output = F::Output> + 's> {
		#![allow(clippy::type_complexity)]
		AnyOf::new(
			self.sim,
			Cons(
				Term(
					Job::build()
						.with_actions(self.fut)
						.with_location(self.loc)
						.with_finalizer(NotifyOnExit::<C>::new())
						.finish()
						.into_inner(),
				),
				Job::build()
					.with_actions(fut)
					.with_finalizer(NotifyOnExit::<C>::new())
					.finish()
					.into_inner(),
			),
		)
	}

	/// Converts the `Fork` into a [`AllOf`]-type of join.
	#[track_caller]
	pub fn and<G: IntoFuture>(
		self,
		fut: G,
	) -> AllOf<'s, C, impl Joinable<C, Output = ((F::Output,), G::Output)> + use<G, C, F>> {
		#![allow(clippy::type_complexity)]
		AllOf::new(
			self.sim,
			Cons(
				Term(
					Job::build()
						.with_actions(self.fut)
						.with_location(self.loc)
						.with_finalizer(NotifyQuorum::<C>::new())
						.finish()
						.into_inner(),
				),
				Job::build()
					.with_actions(fut)
					.with_finalizer(NotifyQuorum::<C>::new())
					.finish()
					.into_inner(),
			),
		)
	}
}

impl<C: ?Sized + Config, F: IntoFuture> IntoFuture for Fork<'_, C, F> {
	type Output = F::Output;
	type IntoFuture = F::IntoFuture;

	fn into_future(self) -> Self::IntoFuture {
		self.fut.into_future()
	}
}

/* ******************************************************* Nested Tuple Lists */

/// A heterogeneous expression list of length one, containing only a head
/// element, but no tail.
#[pin_project::pin_project]
struct Term<H>(#[pin] H);

/// A heterogeneous expression list, containing a single-element head and a tail
/// list.
#[pin_project::pin_project]
struct Cons<H, T>(#[pin] H, #[pin] T);

/* *************************************************************** ActiveList */

/// Represents a list of heterogeneous expressions that are [Active].
pub trait ActiveList<C: ?Sized + Config> {
	/// The length of the list of active expressions.
	const LEN: usize;

	/// A method to activate the active objects.
	fn activate(self: Pin<&mut Self>, sim: &Sim<C>);
}

impl<C, T> ActiveList<C> for Term<T>
where
	C: ?Sized + Config,
	T: Active<C>,
{
	const LEN: usize = 1;

	#[inline]
	fn activate(self: Pin<&mut Self>, sim: &Sim<C>) {
		let this = self.project();

		// SAFETY: This pretends that we don't keep a reference outside the
		// call, which we absolutely do and even use for later querying of the
		// results. Once used, the outside reference invalidates all references
		// kept by the simulator under stacked borrow rules. Despite this, it is
		// safe (though not pretty), because we only use the outside reference
		// once any or all of the jobs have terminated and discard of the
		// other references immediately after.
		// 7/10 - would not transmute again.
		sim.activate(unsafe {
			core::mem::transmute::<Pin<&mut T>, Pin<&mut Lease<'_, T>>>(this.0)
		});
	}
}

impl<C, L, R> ActiveList<C> for Cons<L, R>
where
	C: ?Sized + Config,
	L: ActiveList<C>,
	R: Active<C>,
{
	const LEN: usize = L::LEN + 1;

	#[inline]
	fn activate(self: Pin<&mut Self>, sim: &Sim<C>) {
		let this = self.project();
		this.0.activate(sim);

		// SAFETY: This pretends that we don't keep a reference outside the
		// call, which we absolutely do and even use for later querying of the
		// results. Once used, the outside reference invalidates all references
		// kept by the simulator under stacked borrow rules. Despite this, it is
		// safe (though not pretty), because we only use the outside reference
		// once any or all of the jobs have terminated and discard of the
		// other references immediately after.
		// 7/10 - would not transmute again.
		sim.activate(unsafe {
			core::mem::transmute::<Pin<&mut R>, Pin<&mut Lease<'_, R>>>(this.1)
		});
	}
}