odem_rs_sync/fork/

any_of.rs

use super::{ActiveList, Cons, Term};

use odem_rs_core::{
	ExitStatus, Puck,
	config::Config,
	continuation::Puck as TaskPuck,
	fsm::Brand,
	job::{Job, Settle},
	simulator::{Prec, Sim},
};

use core::{
	future::{Future, IntoFuture},
	pin::Pin,
	task::{Context, Poll},
};

/* *************************************************************** AnyOf Type */

/// A combinator for awaiting the first future to complete.
///
/// This utility resolves when the first future in its collection completes,
/// returning that future's result.
///
/// # Behavior
/// - If multiple futures terminate simultaneously, the one appearing first in
///   the code will take precedence.
/// - Supports chaining with the [`or`] method to combine futures.
///
/// [`or`]: AnyOf::or
#[pin_project::pin_project]
pub struct AnyOf<'s, C: ?Sized + Config, P: Pickable<C>> {
	/// The simulation context before activation and `None` afterward.
	sim: Option<&'s Sim<C>>,
	/// List of futures to pick from.
	#[pin]
	fut: P,
}

impl<'s, C: ?Sized + Config, P: Pickable<C>> AnyOf<'s, C, P> {
	/// Creates a new future combinator containing a set of [Pickable] futures.
	pub(super) const fn new(sim: &'s Sim<C>, fut: P) -> Self {
		AnyOf {
			sim: Some(sim),
			fut,
		}
	}

	/// Adds another future to await to the list.
	#[track_caller]
	pub fn or<F: IntoFuture<Output = P::Output>>(
		self,
		fut: F,
	) -> AnyOf<'s, C, impl Pickable<C, Output = P::Output> + use<F, C, P>> {
		AnyOf {
			sim: self.sim,
			fut: Cons(
				self.fut,
				Job::build()
					.with_actions(fut)
					.with_finalizer(NotifyOnExit::<C>::new())
					.finish()
					.into_inner(),
			),
		}
	}
}

impl<C: ?Sized + Config, P: Pickable<C>> Future for AnyOf<'_, C, P> {
	type Output = P::Output;

	fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
		let mut this = self.project();

		if let Some(sim) = this.sim.take() {
			let active = sim.active();
			let prec = active.prec();

			// distribute the active continuation's puck
			this.fut.as_mut().bind(
				active,
				&mut prec
					.split(P::LEN as _)
					.expect("ran out of bits to support this level of nesting"),
			);

			// activate the list of jobs
			this.fut.activate(sim);
		} else {
			let Some(result) = this.fut.collect() else {
				unreachable!("unexpected spurious reactivation")
			};

			// return the result
			return Poll::Ready(result);
		}

		// await another activation
		Poll::Pending
	}
}

/* ************************************************************* NotifyOnExit */

/// Future adapter that reactivates the parent continuation when the child
/// continuation terminates. Doesn't reactivate the parent if the child is
/// canceled prematurely.
pub struct NotifyOnExit<C: ?Sized + Config> {
	/// Handle to the parent continuation.
	parent: Option<TaskPuck<C>>,
}

impl<C: ?Sized + Config> NotifyOnExit<C> {
	/// Initializes the notifier from a future. The parent continuation is initialized
	/// later during binding to ensure the reactivation of the correct continuation,
	/// even if the job has been moved to a different continuation between creation
	/// and activation.
	pub const fn new() -> Self {
		NotifyOnExit { parent: None }
	}

	/// Method initializing the parent continuation.
	fn bind(&mut self, puck: TaskPuck<C>) {
		self.parent = Some(puck);
	}
}

impl<C: ?Sized + Config, R> Settle<R> for NotifyOnExit<C> {
	fn settle(self, _: &mut R) -> ExitStatus {
		// wake the parent; if it has already been woken up, ignore
		if let Some(mut puck) = self.parent {
			puck.wake().ok();
		}

		// report success
		Ok(odem_rs_core::Success)
	}
}

/* *************************************************************** Pick Trait */

/// Represents a heterogeneous list of [`Job`], of which the earliest one in
/// terms of model time and (secondary) ordering is to be awaited.
pub trait Pickable<C: ?Sized + Config>: ActiveList<C> {
	/// The shared output type of all contained futures.
	type Output;

	/// Completes initializing the `NotifyOnExit` finisher with the caller's
	/// [`TaskPuck`] and sets [`Prec`]s according to the order of queuing.
	fn bind(self: Pin<&mut Self>, puck: TaskPuck<C>, prec: &mut impl Iterator<Item = Prec>);

	/// Collects the result of the single job that terminated.
	fn collect(self: Pin<&mut Self>) -> Option<Self::Output>;
}

impl<C, F> Pickable<C> for Term<Job<C, F, NotifyOnExit<C>>>
where
	C: ?Sized + Config,
	F: Future,
{
	type Output = F::Output;

	fn bind(self: Pin<&mut Self>, puck: TaskPuck<C>, prec: &mut impl Iterator<Item = Prec>) {
		let this = self.project();
		this.0.brand(|job, once| {
			let born = job.token(once).into_born().unwrap();
			job.set_prec(prec.next().unwrap());
			job.finalizer(&born).bind(puck);
		});
	}

	fn collect(self: Pin<&mut Self>) -> Option<Self::Output> {
		let this = self.project();
		this.0.result().ok()
	}
}

impl<C, L, F> Pickable<C> for Cons<L, Job<C, F, NotifyOnExit<C>>>
where
	C: ?Sized + Config,
	L: Pickable<C, Output = F::Output>,
	F: Future,
{
	type Output = F::Output;

	fn bind(self: Pin<&mut Self>, puck: TaskPuck<C>, prec: &mut impl Iterator<Item = Prec>) {
		let this = self.project();

		this.0.bind(puck.clone(), prec);
		this.1.brand(|job, once| {
			let born = job.token(once).into_born().unwrap();
			job.set_prec(prec.next().unwrap());
			job.finalizer(&born).bind(puck);
		});
	}

	fn collect(self: Pin<&mut Self>) -> Option<Self::Output> {
		let this = self.project();
		this.0.collect().or_else(|| this.1.result().ok())
	}
}