odem_rs_sync/fork/

all_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::Sim,
};

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

/* *************************************************************** AllOf Type */

/// A combinator for awaiting the completion of multiple futures.
///
/// This utility allows you to combine multiple [`IntoFuture`] futures into a
/// single future that resolves only after all input futures have completed.
///
/// # Safety Warning
/// - The ordering of fields in [`AllOf`] must not be changed. Specifically, the
///   list of [`Joinable`] futures must precede the `JoinState`. Incorrect
///   ordering can lead to undefined behavior due to dangling pointers.
#[pin_project::pin_project]
pub struct AllOf<'s, C: ?Sized + Config, J: Joinable<C>> {
	/// Reference to the stored simulation context.
	sim: &'s Sim<C>,
	/// List of futures to join.
	#[pin]
	fut: J,
	/// State used to manage completed futures and signal completion.
	#[pin]
	state: JoinState<C>,
}

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

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

impl<C: ?Sized + Config, J: Joinable<C>> Future for AllOf<'_, C, J>
where
	J::Output: Flatten,
{
	type Output = <J::Output as Flatten>::Flattened;

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

		// project the pin onto the continuations
		let mut continuations = this.fut;

		// match over the state counter
		match this.state.counter.get() {
			// has it been initialized already?
			usize::MAX => {
				let state = this.state.into_ref();
				// initialize with the final number of continuations to complete
				state.counter.set(J::LEN);
				// initialize the calling continuation pin
				state.parent.set(Some(this.sim.active()));

				// SAFETY: the `JoinState` outlives the list of jobs activated
				// here by the ordering in the `AllOf` structure in combination
				// with drop order guarantees.
				// Since the type is pinned, we also know that it will be
				// dropped eventually or the program ends, whichever is first.
				unsafe {
					continuations.as_mut().bind(state);
				}

				// activate all continuations
				continuations.activate(this.sim);
			}
			// have all the sub-continuations completed successfully?
			0 => return Poll::Ready(continuations.collect().flatten()),
			// have we been reawakened spuriously?
			_ => unreachable!("unexpected spurious reactivation"),
		}

		Poll::Pending
	}
}

/// The state shared with the continuations generated by [AllOf].
pub struct JoinState<C: ?Sized + Config> {
	/// Number of continuations that have to be completed before the [AllOf]-future
	/// can return.
	counter: Cell<usize>,
	/// Puck of the calling [AllOf]-future. Is reactivated only by the continuation
	/// reducing the [counter] to 0.
	///
	/// [counter]: #structfield.counter
	parent: Cell<Option<TaskPuck<C>>>,
}

impl<C: ?Sized + Config> JoinState<C> {
	/// Creates a new, default `JoinState`.
	const fn new() -> Self {
		Self {
			counter: Cell::new(usize::MAX),
			parent: Cell::new(None),
		}
	}
}

/* ************************************************************* NotifyQuorum */

/// Future adapter that reactivates a given continuation when a quorum in regard to
/// the number of completed sub-continuations is reached.
pub(super) struct NotifyQuorum<C: ?Sized + Config> {
	/// Reference to the shared [`JoinState`] of the caller.
	shared: Option<NonNull<JoinState<C>>>,
}

impl<C: ?Sized + Config> NotifyQuorum<C> {
	/// Initializes the notifier from a future. Parent continuation and quorum are
	/// 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(super) const fn new() -> Self {
		NotifyQuorum { shared: None }
	}

	/// Method initializing the parent continuation and the shared quorum countdown.
	///
	/// # Safety
	/// The caller is responsible that the shared [`JoinState`] outlives the
	/// [`Job`] using this finisher.
	unsafe fn bind(&mut self, state: Pin<&JoinState<C>>) {
		self.shared = Some(NonNull::from(&*state));
	}
}

impl<C: ?Sized + Config, R> Settle<R> for NotifyQuorum<C> {
	fn settle(self, _: &mut R) -> ExitStatus {
		// SAFETY: the bind-method guaranteed that the `JoinState` outlives the
		// job that just completed.
		let state = unsafe { self.shared.unwrap().as_ref() };

		// decrease the number of expected terminations
		let n = state.counter.get() - 1;

		// if it is reduced to 0, activate the parent
		if n == 0 {
			if let Some(mut puck) = state.parent.take() {
				puck.wake().ok();
			}
		}

		// update the counter
		state.counter.set(n);

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

/* *************************************************************** Join Trait */

/// Represents a heterogeneous list of [`Job`] to be awaited.
pub trait Joinable<C: ?Sized + Config>: ActiveList<C> {
	/// The result tuple type once all jobs complete.
	type Output;

	/// Completes initializing the `NotifyQuorum` finisher with the shared
	/// `JoinState`.
	///
	/// # Safety
	/// The caller has to guarantee that the shared `JoinState` outlives all
	/// `Job`s contained in this list.
	unsafe fn bind(self: Pin<&mut Self>, state: Pin<&JoinState<C>>);

	/// Collects the result list containing the outputs of the futures.
	///
	/// This method panics if not all jobs have terminated.
	fn collect(self: Pin<&mut Self>) -> Self::Output;
}

impl<C, F> Joinable<C> for Term<Job<C, F, NotifyQuorum<C>>>
where
	C: ?Sized + Config,
	F: Future,
{
	type Output = (F::Output,);

	unsafe fn bind(self: Pin<&mut Self>, state: Pin<&JoinState<C>>) {
		let this = self.project();
		this.0.brand(|job, once| {
			let born = job.token(once).into_born().unwrap();
			unsafe {
				job.finalizer(&born).bind(state);
			}
		});
	}

	fn collect(self: Pin<&mut Self>) -> Self::Output {
		let this = self.project();
		(this.0.result().expect("activation without result"),)
	}
}

impl<C, L, F> Joinable<C> for Cons<L, Job<C, F, NotifyQuorum<C>>>
where
	C: ?Sized + Config,
	L: Joinable<C>,
	F: Future,
{
	type Output = (L::Output, F::Output);

	unsafe fn bind(self: Pin<&mut Self>, state: Pin<&JoinState<C>>) {
		let this = self.project();

		unsafe {
			this.0.bind(state);
		}
		this.1.brand(|job, once| {
			let born = job.token(once).into_born().unwrap();
			unsafe {
				job.finalizer(&born).bind(state);
			}
		});
	}

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

		(
			this.0.collect(),
			this.1.result().expect("activation without result"),
		)
	}
}
/* ****************************************************************** Flatten */

/// Helper trait that allows flattening of the recursively nested tuples
/// produced in [Joinable::Output].
///
/// The trait is only implemented for the cases required by [`AllOf`]
/// implemented here, not in general for all kinds of recursive nesting.
pub trait Flatten {
	/// Result type after the flattening.
	type Flattened;

	/// Method to perform the flattening.
	fn flatten(self) -> Self::Flattened;
}

/// Helper macro to implement the Flatten-trait generically.
macro_rules! impl_flatten {
	(@name => $T:tt) => { $T };

	(@name $HID:ident $($TID:ident)* =>) => {
		impl_flatten!(@name $($TID)* => ($HID,))
	};

	(@name $HID:ident $($TID:ident)* => $T:tt) => {
		impl_flatten!(@name $($TID)* => ($T, $HID))
	};

	(@leaf $HID:ident) => {
		impl<$HID> Flatten for ($HID,) {
			type Flattened = $HID;

			fn flatten(self) -> Self::Flattened { self.0 }
		}
	};

	(@leaf $($FID:ident)+) => {
		impl<$($FID,)*> Flatten for impl_flatten!(@name $($FID)* =>) {
			type Flattened = ($($FID,)*);

			fn flatten(self) -> Self::Flattened {
				#[allow(non_snake_case)]
				let impl_flatten!(@name $($FID)* =>) = self;
				($($FID,)*)
			}
		}
	};

	(@node $($FID:ident)* .) => {
		impl<H: Flatten, $($FID),*> Flatten for impl_flatten!(@name $($FID)* => H) {
			type Flattened = (H::Flattened, $($FID),*);

			fn flatten(self) -> Self::Flattened {
				#[allow(non_snake_case)]
				let impl_flatten!(@name $($FID)* => H) = self;
				(H.flatten(), $($FID,)*)
			}
		}
	};

	(@node $($HID:ident)* . $MID:ident $($TID:ident)*) => {
		impl_flatten!(@leaf $($HID)* $MID);
		impl_flatten!(@node $($HID)* $MID . $($TID)*);
	};

	($($FID:ident),* $(,)?) => {
		impl_flatten!(@node . $($FID)*);
	};
}

// specialize tuple-flattening for up to 12-tuples before recursively returning
// the remaining tuples partially flattened; this should be enough for most
// cases and also allows debug printing the tuple's contents, which is only
// implemented up to tuple-size 12
impl_flatten!(T1, T2, T3, T4, T5, T6, T7, T8, T9, TA, TB);