odem-rs-core 0.3.0

use crate::{
	ExitStatus,
	config::Config,
	continuation::{Continuation, Label, Share, erased::State, token},
	error::NotInit,
	fsm::*,
	ptr::AsIrc,
	ptr::Irc,
	simulator::{Prec, Sim},
};

use core::{
	any::Any,
	fmt,
	future::IntoFuture,
	mem::ManuallyDrop,
	panic::Location,
	ptr::NonNull,
	task::{RawWaker, RawWakerVTable},
};

/* *********************************************************** Continuation Reference */

/// Reference-counted wrapper around a pinned continuation.
///
/// Instances of this type may be cloned and moved around the simulation model.
/// They allow access to the shared agent state and implement the abstract
/// [`Waker`](core::task::Waker) interface.
///
/// The reference counter does not prolong the lifetime of a [`Continuation`]
/// but instead serves as an indicator that dropping the referenced continuation
/// would lead to dangling references if its reference counter is greater than
/// zero. Thus, the actual mechanics of this type are more akin to a checksum
/// than a reference-counted type since it falls upon the user to drop their
/// `Puck`s before the referenced continuation is dropped.
///
/// This design preserves memory-safety for continuations with erased lifetimes
/// at the cost of a reduction in user-friendliness since bugs only surface
/// during destruction of continuations, making them potentially hard to track.
///
/// However, it also allows for `no_std` environments without dynamic memory
/// allocation to run simulations as it is not necessary to keep continuations
/// on the heap.
pub struct Puck<C: ?Sized + Config>(Irc<Continuation<'static, C>>);

impl<C: ?Sized + Config> Puck<C> {
	/// Creates a new [`Puck`] from a [`Continuation`] that has been initialized
	/// through binding.
	pub(crate) fn new<'b, T>(task: Irc<Continuation<'b, C>>, _: &T) -> Self
	where
		T: Into<token::Init<'b>>,
	{
		Puck(Irc::map(task, Continuation::detach))
	}

	/// Creates a new [`Puck`] from a [`Continuation`] and returns `Err` if it
	/// hadn't been initialized.
	pub(crate) fn checked(task: Irc<Continuation<'static, C>>) -> Result<Self, NotInit> {
		let state = task.state().erased();
		state.is_init().then(|| Puck(task)).ok_or(NotInit(state))
	}

	/// Extracts the reference-counted continuation from the puck.
	pub(crate) fn into_inner(self) -> Irc<Continuation<'static, C>> {
		self.0
	}

	/// Returns a reference to the [shared context](Share) bound to this
	/// continuation.
	pub(crate) fn share(&self) -> &Share<C> {
		// SAFETY: Pucks can only be created during initialization, after
		// which they have shared data bound to them
		unsafe { self.0.share().unwrap_unchecked() }
	}

	/// Returns whether this continuation points to a given [`Continuation`].
	pub(crate) fn is_same(&self, other: &Continuation<'_, C>) -> bool {
		core::ptr::eq(&*self.0, other.detach())
	}
}

// implement the waker interface for continuation pins
impl<C: Config> Puck<C> {
	const VTABLE: RawWakerVTable = RawWakerVTable::new(
		Self::waker_clone,
		Self::waker_wake,
		Self::waker_wake_by_ref,
		Self::waker_drop,
	);

	/// Converts a continuation pin into a [`RawWaker`] without decreasing the
	/// reference count of the inner continuation.
	pub(crate) fn into_waker(self) -> RawWaker {
		// initialize the raw waker
		RawWaker::new(Irc::into_raw(self.0).cast().as_ptr(), &Self::VTABLE)
	}

	unsafe fn waker_clone(task: *const ()) -> RawWaker {
		// restore the puck from the raw pointer
		let task_ptr = NonNull::new(task as *mut Continuation<'static, C>).unwrap();
		let task = unsafe { Irc::from_raw(task_ptr) };

		// assert that the waker hasn't been called from a different thread
		// SAFETY: the underlying continuation has been activated before
		if !unsafe { task.is_same_thread() } {
			abort();
		}

		let puck = ManuallyDrop::new(Puck(task));

		// create another (identical) raw waker
		ManuallyDrop::into_inner(puck.clone()).into_waker()
	}

	unsafe fn waker_wake(task: *const ()) {
		let task_ptr = NonNull::new(task as *mut Continuation<'static, C>).unwrap();
		let task = unsafe { Irc::from_raw(task_ptr) };

		// assert that the waker hasn't been called from a different thread
		// SAFETY: the underlying continuation has been activated before
		if !unsafe { task.is_same_thread() } {
			abort();
		}

		Continuation::wake(task).ok();
	}

	unsafe fn waker_wake_by_ref(task: *const ()) {
		let task_ptr = NonNull::new(task as *mut Continuation<'static, C>).unwrap();
		let task = ManuallyDrop::new(unsafe { Irc::from_raw(task_ptr) });

		// assert that the waker hasn't been called from a different thread
		// SAFETY: the underlying continuation has been activated before
		if !unsafe { task.is_same_thread() } {
			abort();
		}

		Continuation::wake((*task).clone()).ok();
	}

	unsafe fn waker_drop(task: *const ()) {
		let task_ptr = NonNull::new(task as *mut Continuation<'static, C>).unwrap();
		let task = unsafe { Irc::from_raw(task_ptr) };

		// assert that the waker hasn't been called from a different thread
		// SAFETY: the underlying continuation has been activated before
		if !unsafe { task.is_same_thread() } {
			abort();
		}
	}
}

impl<C: ?Sized + Config> crate::Puck<C> for Puck<C> {
	fn result(&mut self) -> Option<ExitStatus> {
		self.0.result()
	}

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

	fn sim(&self) -> &Sim<C> {
		self.share().sim()
	}

	fn label(&self) -> Label {
		self.share().label()
	}

	fn time(&self) -> Option<C::Time> {
		self.0.time()
	}

	fn rank(&self) -> C::Rank {
		self.share().rank()
	}

	fn prec(&self) -> Prec {
		self.0.prec()
	}

	fn state(&self) -> State {
		self.0.state().borrow().erased()
	}

	fn location(&self) -> &'static Location<'static> {
		self.0.location()
	}

	fn puck(&self) -> Puck<C> {
		self.clone()
	}
}

impl<C: ?Sized + Config> AsRef<Continuation<'static, C>> for Puck<C> {
	fn as_ref(&self) -> &Continuation<'static, C> {
		&self.0
	}
}

impl<C: ?Sized + Config> AsIrc<Continuation<'static, C>> for Puck<C> {
	fn as_irc(&self) -> Irc<Continuation<'static, C>> {
		self.0.clone()
	}
}

impl<C: ?Sized + Config> IntoFuture for Puck<C> {
	type Output = ExitStatus;
	type IntoFuture = crate::ops::Join<C, Self>;

	fn into_future(self) -> Self::IntoFuture {
		crate::ops::join(self)
	}
}

impl<C: ?Sized + Config> fmt::Debug for Puck<C> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		self.0.fmt(f)
	}
}

impl<C: ?Sized + Config> Clone for Puck<C> {
	fn clone(&self) -> Self {
		Puck(self.0.clone())
	}
}

impl<C: ?Sized + Config> PartialEq for Puck<C> {
	fn eq(&self, other: &Self) -> bool {
		self.is_same(&other.0)
	}
}

impl<C: ?Sized + Config> Eq for Puck<C> {}

extern "C" fn abort() -> ! {
	panic!("attempted to use a `Waker` from a different thread");
}

#[cfg(test)]
mod tests {
	use crate::{
		job::Job,
		ops::{defer, waker},
		simulator::{Sim, simulation},
	};
	use std::{pin::pin, sync::mpsc::channel, thread};

	#[test]
	#[ignore = "requires abort test runner"]
	fn move_waker_across_thread() {
		simulation(async |sim: &Sim| {
			// Create a Waker and move it across thread-boundaries.
			let waker = waker().await;

			let job = pin!(Job::new(async move {
				let handle = thread::spawn(move || {
					// Receive the waker in another thread and use it.
					waker.wake();
				});

				// The waker should lead to a panic in the other thread.
				assert!(handle.join().is_err());
			}));

			// Wait for the Job to complete.
			sim.activate(job).await;
		})
		.expect("unexpected deadlock");
	}

	#[test]
	#[ignore = "requires abort test runner"]
	#[should_panic = "attempted to use a `Waker` from a different thread"]
	fn move_waker_across_thread_channel() {
		let (sx, rx) = channel();

		thread::spawn(move || {
			simulation(async |_: &Sim| {
				// Create a Waker and move it across thread-boundaries.
				let waker = waker().await;
				sx.send(waker).unwrap();
				loop {
					defer().await;
				}
			})
			.expect("unexpected deadlock");
		});

		// The waker-call should lead to a panic.
		let waker = rx.recv().unwrap();
		waker.wake();
	}

	#[test]
	#[ignore = "requires abort test runner"]
	fn move_waker_across_thread_race() {
		let (sx, rx) = channel();

		thread::spawn(move || {
			simulation(async |_: &Sim| {
				// Create a Waker and move it across thread-boundaries.
				let waker = waker().await;
				sx.send(waker).unwrap();
			})
			.expect("unexpected deadlock");
		});

		// The waker-call should lead to a panic.
		let waker = rx.recv().unwrap();
		waker.wake();
	}
}