use core::{
cell::{BorrowMutError, Cell, Ref, RefCell},
marker::PhantomData,
ops::Deref,
pin::Pin,
};
pub trait FSM {
type Token<'b>;
type Erased;
fn token<'b>(this: &Self, once: Ephemeral<'b>, _: Private) -> Self::Token<'b>;
fn erased(&self) -> Self::Erased;
fn label(&self) -> &'static str;
}
pub trait Stateful {
type Brand;
unsafe fn enter(&self);
unsafe fn leave(&self);
}
pub trait Rebrand<'brand>: Stateful<Brand = &'brand ()> {
type Kind<'a>: Rebrand<'a, Kind<'brand> = Self>;
fn identity_ref(&self) -> &Self::Kind<'brand> {
unsafe { &*(self as *const Self as *const Self::Kind<'brand>) }
}
fn identity_mut(&mut self) -> &mut Self::Kind<'brand> {
unsafe { &mut *(self as *mut Self as *mut Self::Kind<'brand>) }
}
}
pub trait Brand<F, R> {
fn brand(self, f: F) -> R;
}
pub trait Debrand<'b, F, R> {
fn debrand(self, once: impl Into<Ephemeral<'b>>, f: F) -> (Ephemeral<'b>, R);
}
pub trait Transition<'b, U, V> {
type Data;
fn transition(this: &mut Self, curr: U, data: Self::Data, _: Private) -> V;
}
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Ephemeral<'brand>(PhantomData<fn(&'brand ()) -> &'brand ()>);
pub struct Private(());
pub struct StateMachine<'b, T: ?Sized> {
_mark: PhantomData<fn(&'b ()) -> &'b ()>,
token: Cell<u32>,
state: RefCell<T>,
}
impl<T> StateMachine<'static, T> {
pub const fn new(inner: T) -> Self {
Self {
_mark: PhantomData,
token: Cell::new(0),
state: RefCell::new(inner),
}
}
}
impl<'b, T: FSM> StateMachine<'b, T> {
pub fn token(&self, once: Ephemeral<'b>) -> T::Token<'b> {
T::token(&*self.borrow(), once, Private(()))
}
pub fn erased(&self) -> T::Erased {
T::erased(&*self.borrow())
}
pub fn label(&self) -> &'static str {
T::label(&*self.borrow())
}
pub fn transition<U, V>(&self, curr: U, data: T::Data) -> V
where
T: Transition<'b, U, V>,
{
match self.try_transition(curr, data) {
Ok(v) => v,
Err(err) => panic!(
"cannot transition from '{}' to '{}': {err}",
self.label(),
core::any::type_name::<V>(),
),
}
}
pub fn try_transition<U, V>(&self, curr: U, data: T::Data) -> Result<V, TransitionError<U>>
where
T: Transition<'b, U, V>,
{
match self.token.get() {
0 => unreachable!("transitioning without tokens should be impossible"),
1 => match self.state.try_borrow_mut() {
Ok(mut state) => Ok(T::transition(&mut *state, curr, data, Private(()))),
Err(err) => Err(TransitionError::Borrow(curr, err)),
},
n => Err(TransitionError::Invalidation(curr, n)),
}
}
pub fn update<U>(&self, curr: U, data: T::Data) -> U
where
T: Transition<'b, U, U>,
{
T::transition(&mut *self.state.borrow_mut(), curr, data, Private(()))
}
#[track_caller]
pub fn borrow(&self) -> Ref<'_, T> {
self.state.borrow()
}
}
impl<T: Default> Default for StateMachine<'static, T> {
fn default() -> Self {
Self::new(T::default())
}
}
impl<'b, T> Stateful for StateMachine<'b, T> {
type Brand = &'b ();
#[inline]
unsafe fn enter(&self) {
self.token.set(self.token.get() + 1);
}
#[inline]
unsafe fn leave(&self) {
self.token.set(self.token.get() - 1);
}
}
impl<'b, T> Rebrand<'b> for StateMachine<'b, T> {
type Kind<'a> = StateMachine<'a, T>;
}
#[derive(thiserror::Error)]
pub enum TransitionError<T> {
#[error("cannot transition while {1} references exist")]
Invalidation(T, u32),
#[error("cannot transition while state is borrowed")]
Borrow(T, BorrowMutError),
}
struct Guard<T: Stateful, const E: bool = true>(T);
impl<T: Stateful> Guard<T, true> {
fn enter(value: T) -> Self {
unsafe {
value.enter();
}
Self(value)
}
}
impl<T: Stateful> Guard<T, false> {
fn leave(value: T) -> Self {
unsafe {
value.leave();
}
Self(value)
}
}
impl<T: Stateful, const E: bool> Drop for Guard<T, E> {
fn drop(&mut self) {
if E {
unsafe {
self.0.leave();
}
} else {
unsafe {
self.0.enter();
}
}
}
}
impl<T: Deref<Target: Stateful>> Stateful for T {
type Brand = <T::Target as Stateful>::Brand;
unsafe fn enter(&self) {
unsafe {
(**self).enter();
}
}
unsafe fn leave(&self) {
unsafe {
(**self).leave();
}
}
}
impl<'b, T, F, R> Brand<F, R> for &T
where
T: Rebrand<'b>,
F: for<'a> FnOnce(&T::Kind<'a>, Ephemeral<'a>) -> R,
{
fn brand(self, f: F) -> R {
f(Guard::enter(self.identity_ref()).0, Ephemeral(PhantomData))
}
}
impl<'b, T, F, R> Brand<F, R> for &mut T
where
T: Rebrand<'b>,
F: for<'a> FnOnce(&mut T::Kind<'a>, Ephemeral<'a>) -> R,
{
fn brand(self, f: F) -> R {
f(Guard::enter(self.identity_mut()).0, Ephemeral(PhantomData))
}
}
impl<'b, T, F, R> Brand<F, R> for Pin<&T>
where
T: Rebrand<'b>,
F: for<'a> FnOnce(Pin<&T::Kind<'a>>, Ephemeral<'a>) -> R,
{
fn brand(self, f: F) -> R {
let this = unsafe { Pin::map_unchecked(self, T::identity_ref) };
f(Guard::enter(this).0.as_ref(), Ephemeral(PhantomData))
}
}
impl<'b, T, F, R> Brand<F, R> for Pin<&mut T>
where
T: Rebrand<'b>,
F: for<'a> FnOnce(Pin<&mut T::Kind<'a>>, Ephemeral<'a>) -> R,
{
fn brand(self, f: F) -> R {
let this = unsafe { Pin::map_unchecked_mut(self, T::identity_mut) };
f(Guard::enter(this).0.as_mut(), Ephemeral(PhantomData))
}
}
impl<'b, T, F, R> Brand<F, R> for crate::ptr::Irc<T>
where
T: crate::ptr::IntrusivelyCounted + Rebrand<'b, Kind<'b>: crate::ptr::IntrusivelyCounted>,
F: for<'a> FnOnce(&crate::ptr::Irc<T::Kind<'a>>, Ephemeral<'a>) -> R,
{
fn brand(self, f: F) -> R {
let this = crate::ptr::Irc::map(self, T::identity_ref);
f(&Guard::enter(this).0, Ephemeral(PhantomData))
}
}
impl<'b, T, F, R> Debrand<'b, F, R> for &T
where
T: Rebrand<'b>,
F: FnOnce(&T) -> R,
{
fn debrand(self, once: impl Into<Ephemeral<'b>>, f: F) -> (Ephemeral<'b>, R) {
let guard = Guard::leave(self);
(once.into(), f(guard.0))
}
}
impl<'b, T, F, R> Debrand<'b, F, R> for Pin<&T>
where
T: Rebrand<'b>,
F: FnOnce(Pin<&T>) -> R,
{
fn debrand(self, once: impl Into<Ephemeral<'b>>, f: F) -> (Ephemeral<'b>, R) {
let guard = Guard::leave(self);
(once.into(), f(guard.0))
}
}
impl<'b, T, F, R> Debrand<'b, F, R> for &mut T
where
T: Rebrand<'b>,
F: FnOnce(&mut T) -> R,
{
fn debrand(self, once: impl Into<Ephemeral<'b>>, f: F) -> (Ephemeral<'b>, R) {
let guard = Guard::leave(self);
(once.into(), f(guard.0))
}
}
impl<'b, T, F, R> Debrand<'b, F, R> for Pin<&mut T>
where
T: Rebrand<'b>,
F: FnOnce(Pin<&mut T>) -> R,
{
fn debrand(self, once: impl Into<Ephemeral<'b>>, f: F) -> (Ephemeral<'b>, R) {
let mut guard = Guard::leave(self);
(once.into(), f(guard.0.as_mut()))
}
}
#[doc(hidden)]
#[macro_export]
macro_rules! __fsm {
(
$(#[$DOC:meta])*
$V:vis enum $E:ident $(<$C:ident: $W:ty>)? {
$($(#[$SDOC:meta])* $S:ident $(($T:ty))? -> {$($D:ident),* $(,)?}),* $(,)?
}
$(
$(#[$MDOC:meta])*
$MV:vis $MS:ident $(: $MB:ident)? = {$($MT:ident),* $(,)?};
)*
) => {paste::paste! {
#[cfg_attr(doc, aquamarine::aquamarine)]
$(#[$DOC])*
#[doc = "```mermaid"]
#[doc = "---"]
#[doc = "title: " $E " Transition Diagram"]
#[doc = "---"]
#[doc = "flowchart LR"]
#[doc = $($("\t" $S "{{" $S "}}-->" $D "{{" $D "}}\n")*)*]
#[doc = "```"]
$V enum $E $(<$C: ?Sized + $W>)* {
$(
$(#[$SDOC])*
$S $(($T))*
),*
}
impl<$($C: ?Sized + $W)*> $crate::fsm::FSM for $E $(<$C>)* {
type Token<'brand> = token::$E<'brand>;
type Erased = erased::$E;
fn token<'brand>(this: &Self, once: Ephemeral<'brand>, _: Private) -> Self::Token<'brand> {
match this {$(
Self::$S {..} => token::$E::$S(token::$S(once))
),*}
}
fn erased(&self) -> Self::Erased {
match self {$(
Self::$S {..} => erased::$E::$S
),*}
}
fn label(&self) -> &'static str {
match self {$(
Self::$S {..} => stringify!($S)
),*}
}
}
#[allow(dead_code)]
impl<$($C: ?Sized + $W)*> $E $(<$C>)* {
$(
#[doc = "Returns `true` if the current state is "]
#[doc = "[`Self::" $S "`] and `false` otherwise."]
pub const fn [<is_ $S:lower>](&self) -> bool {
matches!(self, Self::$S {..})
}
)*
$(
#[doc = "Returns `true` if the current state is "]
#[doc = $("[`Self::" $MT "`]")" or "* " and `false` otherwise."]
pub const fn [<is_ $MS:lower>](&self) -> bool {
$(self.[<is_ $MT:lower>]())|*
}
)*
}
#[doc = "A module containing a configuration- and payload-erased "]
#[doc = "version of [`" $E "`]."]
$V mod erased {
#[allow(unused_imports)]
use super::*;
#[cfg_attr(doc, aquamarine::aquamarine)]
$(#[$DOC])*
#[doc = "```mermaid"]
#[doc = "---"]
#[doc = "title: " $E " Transition Diagram"]
#[doc = "---"]
#[doc = "flowchart LR"]
#[doc = $($("\t" $S "{{" $S "}}-->" $D "{{" $D "}}\n")*)*]
#[doc = "```"]
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
pub enum $E {$(
$(#[$SDOC])*
$S
),*}
#[allow(dead_code)]
impl $E {
pub const fn label(&self) -> &'static str {
match self {$(
Self::$S => stringify!($S)
),*}
}
$(
#[doc = "Returns `true` if the current state is [`Self::" $S "`] and `false` otherwise."]
pub const fn [<is_ $S:lower>](&self) -> bool {
matches!(self, Self::$S)
}
)*
$(
#[doc = "Returns `true` if the current state is "]
#[doc = $("[`Self::" $MT "`]")" or "* " and `false` otherwise."]
pub const fn [<is_ $MS:lower>](&self) -> bool {
match self {
$(Self::$MT)|* => true,
_ => false
}
}
)*
}
}
#[doc = "A module containing weightless, branded 👻 token types that "]
#[doc = "are used as testimony that equally branded [`" $E "`] are in "]
#[doc = "the corresponding state (`" $($S)"` or `"* "`)."]
$V mod token {
use $crate::fsm::{Ephemeral, Transition};
use super::*;
use core::fmt;
#[doc = "An enumeration of token-witnesses for the states "]
#[doc = $("[`" $S "`](super::" $E "::" $S ")")", "* " and the meta-states "]
#[doc = $("[`" $MS "`]")", "* "."]
#[derive(Debug, Eq, PartialEq, Hash)]
pub enum $E<'brand> {$(
#[doc = "[`" $S "`] state containing the token as a witness."]
#[doc = ""]
#[doc = "[`" $S "`]: super::" $E "::" $S]
$S($S<'brand>)
),*}
#[allow(dead_code)]
impl<'brand> $E<'brand> {
pub const fn erased(&self) -> erased::$E {
match self {$(
Self::$S(_) => erased::$E::$S
),*}
}
$(
#[doc = "Attempts to convert the token referencing the [`" $E "`] "]
#[doc = "into a token referencing the more specialized [`" $S "`]."]
#[doc = ""]
#[doc = "[`" $E "`]: super::" $E]
#[doc = "[`" $S "`]: super::" $E "::" $S]
pub const fn [<as_ $S:lower>](&self) -> Option<&$S<'brand>> {
match self {
Self::$S(token) => Some(token),
_ => None
}
}
#[doc = "Attempts to convert the token for [`" $E "`] "]
#[doc = "into a token for the more specialized [`" $S "`]."]
#[doc = ""]
#[doc = "[`" $E "`]: super::" $E]
#[doc = "[`" $S "`]: super::" $E "::" $S]
pub const fn [<into_ $S:lower>](self) -> Result<$S<'brand>, error::[<Not $S>]> {
match self {
Self::$S(token) => Ok(token),
_ => Err(error::[<Not $S>](self.erased()))
}
}
)*
$(
#[doc = "Attempts to convert the token referencing the "]
#[doc = "[`" $E "`] into a token referencing state "]
#[doc = $($MT)" or "* "."]
#[doc = ""]
#[doc = "[`" $E "`]: super::" $E]
pub const fn [<as_ $MS:lower>](&self) -> Option<&$MS<'brand>> {
match self {
$(Self::$MT(token) => Some(unsafe { &*(token as *const $MT<'brand> as *const $MS<'brand>) }),)*
_ => None
}
}
#[doc = "Attempts to convert the token for [`" $E "`] "]
#[doc = "into a token for the more specialized "$($MT)" or "* "."]
#[doc = ""]
#[doc = "[`" $E "`]: super::" $E]
pub const fn [<into_ $MS:lower>](self) -> Result<$MS<'brand>, error::[<Not $MS>]> {
match self {
$(Self::$MT(token) => Ok($MS(token.0)),)*
_ => Err(error::[<Not $MS>](self.erased()))
}
}
)*
}
impl<'brand> From<$E<'brand>> for Ephemeral<'brand> {
fn from(value: $E<'brand>) -> Self {
match value {
$($E::$S(token) => token.0),*
}
}
}
$(
#[doc = "Token testifying that the equally branded "]
#[doc = "[`" $E "`] is [`" $S "`]."]
#[doc = ""]
#[doc = "[`" $S "`]: super::" $E "::" $S]
#[derive(PartialOrd, PartialEq, Ord, Eq, Hash)]
#[repr(transparent)]
$V struct $S<'brand>(pub(super) Ephemeral<'brand>);
impl<'brand> $S<'brand> {
pub const unsafe fn new(once: Ephemeral<'brand>) -> Self {
Self(once)
}
}
impl<'brand> From<$S<'brand>> for Ephemeral<'brand> {
fn from(value: $S<'brand>) -> Self {
value.0
}
}
impl fmt::Debug for $S<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, concat!(stringify!($E), "(", stringify!($S), ")"))
}
}
)*
$crate::fsm::fsm! {
@gen_transitions ($E) ($($W)*) ($($S $(($T))* => $($D)*),*)
}
$(
$(#[$MDOC])*
#[repr(transparent)]
$MV struct $MS<'brand>(Ephemeral<'brand>);
#[allow(dead_code)]
impl<'brand> $MS<'brand> {
pub const unsafe fn new(once: Ephemeral<'brand>) -> Self {
Self(once)
}
$(
#[doc = "Re-casts a reference to the specialized "]
#[doc = "[`" $MS "`]-token into a token of the generalized "]
#[doc = "[`" $MB "`] kind."]
pub const fn [<as_ $MB:lower>](&self) -> &$MB<'brand> {
unsafe { &*(self as *const $MS<'brand> as *const $MB<'brand>) }
}
#[doc = "Converts the specialized [`" $MS "`]-token into "]
#[doc = "a token of the generalized [`" $MB "`] kind."]
pub const fn [<into_ $MB:lower>](self) -> $MB<'brand> {
$MB(self.0)
}
)*
}
impl<'brand> From<$MS<'brand>> for Ephemeral<'brand> {
fn from(value: $MS<'brand>) -> Self {
value.0
}
}
impl fmt::Debug for $MS<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, concat!("Token(", stringify!($MS), ")"))
}
}
$(
#[allow(dead_code)]
impl<'brand> $MT<'brand> {
#[doc = "Re-casts a reference to the specialized "]
#[doc = "[`" $MT "`]-token into a token of the generalized "]
#[doc = "[`" $MS "`] kind."]
pub const fn [<as_ $MS:lower>](&self) -> &$MS<'brand> {
unsafe { &*(self as *const $MT<'brand> as *const $MS<'brand>) }
}
#[doc = "Converts the specialized [`" $MT "`]-token into "]
#[doc = "a token of the generalized [`" $MS "`] kind."]
pub const fn [<into_ $MS:lower>](self) -> $MS<'brand> {
$MS(self.0)
}
}
impl<'brand> From<$MT<'brand>> for $MS<'brand> {
fn from(value: $MT<'brand>) -> Self { value.[<into_ $MS:lower>]() }
}
)*
)*
}
$V mod error {
use core::fmt;
$(
#[doc = "Error type indicating that the current [`" $E "`] was not [`" $S "`]."]
#[doc = ""]
#[doc = "[`" $S "`]: super::" $E "::" $S]
#[doc = "[`" $E "`]: super::" $E]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct [<Not $S>](pub super::erased::$E);
impl fmt::Display for [<Not $S>] {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
concat!("<", stringify!($S), "> expected but <{}> found instead"),
self.0.label()
)
}
}
impl core::error::Error for [<Not $S>] {}
)*
$(
#[doc = "Error type indicating that the current [`" $E "`] was "]
#[doc = "none of " $("[`" $MT "`](super::" $E "::" $MT ")")" or "* "."]
#[doc = ""]
#[doc = "[`" $E "`]: super::" $E]
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct [<Not $MS>](pub super::erased::$E);
impl fmt::Display for [<Not $MS>] {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
concat!("either " $(, "'", stringify!($MT), "'",)" or "* " expected but state '{}' found instead"),
self.0.label()
)
}
}
impl core::error::Error for [<Not $MS>] {}
)*
}
}};
(@gen_transitions ($E:ident) ($W:ty) ($($S:ident $(($T:ty))? => $($D:ident)*),*)) => {paste::paste! {$(
impl<'b, C: ?Sized + $W> Transition<'b,$S<'b>, $S<'b>> for super::$E<C> {
#[allow(unused_parens)]
type Data = ($($T)*);
#[doc = "Updates the state of [`" $S "`],"]
#[doc = "requiring a " $S "-token and returning it."]
fn transition(this: &mut Self, token: $S<'b>, _data: Self::Data, _: $crate::fsm::Private) -> $S<'b> {
*this = super::$E::$S $(({
let _: $T;
_data
}))?;
token
}
}
$(
impl<'b, C: ?Sized + $W> Transition<'b,$S<'b>, $D<'b>> for super::$E<C> {
type Data = <Self as Transition<'b,$D<'b>,$D<'b>>>::Data;
#[doc = "Performs the transition between [`" $S "`] and [`" $D "`],"]
#[doc = "requiring a " $S "-token and returning the resulting "]
#[doc = $D "-token."]
#[cfg_attr(feature = "debug-tracing", track_caller)]
fn transition(this: &mut Self, token: $S<'b>, data: Self::Data, p: $crate::fsm::Private) -> $D<'b> {
#[cfg(feature = "debug-tracing")]
::tracing::trace!(from = %stringify!($S), into = %stringify!($D), "transition");
Self::transition(this, $D(token.0), data, p)
}
}
)*
)*}};
(@gen_transitions ($E:ident) () ($($S:ident $(($T:ty))? => $($D:ident)*),*)) => {paste::paste! {$(
impl<'b> Transition<'b, $S<'b>, $S<'b>> for super::$E {
#[allow(unused_parens)]
type Data = ($($T)*);
#[doc = "Updates the state of [`" $S "`],"]
#[doc = "requiring a " $S "-token and returning it."]
fn transition(this: &mut Self, token: $S<'b>, _data: Self::Data, _: $crate::fsm::Private) -> $S<'b> {
*this = super::$E::$S $(({
let _: $T;
_data
}))?;
token
}
}
$(
impl<'b> Transition<'b, $S<'b>, $D<'b>> for super::$E {
type Data = <Self as Transition<'b,$D<'b>,$D<'b>>>::Data;
#[doc = "Performs the transition between [`" $S "`] and [`" $D "`],"]
#[doc = "requiring a " $S "-token and returning the resulting "]
#[doc = $D "-token."]
#[cfg_attr(feature = "debug-tracing", track_caller)]
fn transition(this: &mut Self, token: $S<'b>, data: Self::Data, p: $crate::fsm::Private) -> $D<'b> {
#[cfg(feature = "debug-tracing")]
::tracing::trace!(from = %stringify!($S), into = %stringify!($D), "transition");
Self::transition(this, $D(token.0), data, p)
}
}
)*
)*}};
}
#[doc(inline)]
pub use __fsm as fsm;
#[cfg(test)]
mod tests {
use super::*;
fsm! {
#[allow(dead_code)]
pub enum States {
Idle -> { Busy },
Busy -> { Idle, Done },
Done(bool) -> {}
}
}
#[test]
fn simple_sm() {
let m1 = StateMachine::new(States::Idle);
let _ = m1.brand(|sm, once| {
let idle = sm.token(once).into_idle().unwrap();
let busy: token::Busy<'_> = sm.transition(idle, ());
let _: token::Done<'_> = sm.transition(busy, true);
m1.borrow()
});
}
}