odem-rs-core 0.1.0

use crate::{
	ExitStatus,
	config::Config,
	continuation::Share,
	continuation::{Continuation, Label, erased::State, token},
	error::{NotIdle, 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)
	}

	/// Reconstructs the original [`Puck`] from the raw pointer.
	///
	/// This function also ensures that the call didn't come from a different
	/// thread than the one that the executor runs in.
	///
	/// # Safety
	/// The caller has to ensure that the provided raw pointer was the one
	/// originally used in the call to [`Self::into_waker`].
	unsafe fn restore(task_ptr: NonNull<Continuation<'static, C>>) -> Self {
		// reconstruct the pointer to the continuation
		let task = unsafe { task_ptr.as_ref() };

		// determine whether this thread is different from the one that the
		// executor is running in
		let same_thread = task.brand(|task, once| {
			// SAFETY: we originally used a puck to create the waker
			let init = unsafe { task.token(once).into_init().unwrap_unchecked() };

			task.is_same_thread(&init)
		});

		// assert that the waker hasn't been called from a different thread
		assert!(
			same_thread,
			"simulation-associated waker may not be moved across thread boundaries"
		);

		// now it's safe to reconstruct the puck
		Puck(unsafe { Irc::from_raw(task_ptr) })
	}

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

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

	unsafe fn waker_wake(task: *const ()) {
		unsafe {
			// restore the puck and reactivate the underlying continuation
			Continuation::wake(
				Self::restore(NonNull::new(task as *mut Continuation<'static, C>).unwrap()).0,
			)
			.ok();
		}
	}

	unsafe fn waker_wake_by_ref(task: *const ()) {
		unsafe {
			// restore the puck without dropping it later
			let puck = ManuallyDrop::new(Self::restore(
				NonNull::new(task as *mut Continuation<'static, C>).unwrap(),
			));

			// reactivate the underlying continuation
			Continuation::wake(puck.0.clone()).ok();
		}
	}

	unsafe fn waker_drop(task: *const ()) {
		unsafe {
			// restore the puck and drop it immediately
			Self::restore(NonNull::new(task as *mut Continuation<'static, C>).unwrap());
		}
	}
}

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

	fn wake(&mut self) -> Result<(), NotIdle> {
		Continuation::wake(self.0.clone())
	}

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

	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()
	}
}

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> {}