Some approaches to making a state machine in Haskell

StateMachine.hs

module StateMachine where

import Data.Foldable (foldlM)

data Init = Init
  deriving (Eq, Show)
data A = A Int
  deriving (Eq, Show)
data B = B String
  deriving (Eq, Show)
data C = C Int String
  deriving (Eq, Show)

data GotInt = GotInt Int
  deriving (Eq, Show)
data GotString = GotString String
  deriving (Eq, Show)

data State = StateInit Init
           | StateA A
           | StateB B
           | StateC C
  deriving (Eq, Show)

data Event = EventGotInt GotInt
           | EventGotString GotString
  deriving (Eq, Show)

{-
Valid transitions:

Init -> GotInt -> A
Init -> GotString -> B
A -> GotString -> C
B -> GotInt -> C
(C has no valid transitions)
-}

-- First approach:
-- Multi-param type class with single method
-- We have one instance for each valid transition, and one instance to dispatch
-- on our sum types State and Event.
-- We have to wire up the "dispatch function" by hand.
-- The type classes ensure we cannot add in a bad transition.
-- We could, however, forget to use a valid transition, in which case it would
-- fall through to the error case.

class NextState currentState event where
  next :: event -> currentState -> State

-- Specific instances
instance NextState Init GotInt where
  next (GotInt i) Init = StateA (A i)
instance NextState Init GotString where
  next (GotString s) Init = StateB (B s)

instance NextState A GotString where
  next (GotString s) (A i) = StateC (C i s)

instance NextState B GotInt where
  next (GotInt i) (B s) = StateC (C i s)

-- "Dispatch function" that handles our sum types
handleEventMPTC :: Event -> State -> Either String State
handleEventMPTC (EventGotInt e)    (StateInit s) = Right $ next e s
handleEventMPTC (EventGotString e) (StateInit s) = Right $ next e s
handleEventMPTC (EventGotString e) (StateA s)    = Right $ next e s
handleEventMPTC (EventGotInt e)    (StateB s)    = Right $ next e s
handleEventMPTC _                  _             = Left "Invalid Transition"

foldEventsMPTC :: State -> [Event] -> Either String State
foldEventsMPTC = foldlM (flip handleEventMPTC)

-- Second approach:
-- Multi-param type class with single method and catch-all error instance
-- Same as the above, but add a "catch-all" overlappable instance.
-- If there is no specific implementation for a transition, the compiler will
-- use that instance, which always goes to an error state.
-- In this implementation, we wired up a "dispatch instance" by hand.
-- But we do not have to think about what instances exist or not, we
-- just dispatch every possible combination of types. This means we could
-- probably build this automatically using Generics or TemplateHaskell, which
-- would be great and error-proof.
-- This is great if we're folding a list of Events (since there's no way to know
-- at compile-time what that sum type Event will hold at run time, so we have to
-- deal with the error case.
-- Having an error state is inevitable if we want to do something like folding
-- over a list of Events, because there's no way to know at compile-time what
-- each Event will hold at runtime, since Event is a sum type. So this solution
-- is great if what you want to do is fold over a list of Events that you read
-- from user input or from a database.
-- However, we have lost the validation that prevented us from explicitly
-- calling `next (C 0 "What") (GotString "Who")`, because there's now an instance
-- that implements that. Besides, all our transitions return
-- `Either String State`, so all of them could potentially return an error state.
-- If we're doing something different from "folding over a list of Events that
-- we read from user input or from a database", then having those compile-time
-- checks of correctnes and totality may be useful.

class NextStateWithFallbackInstance currentState event where
  nextWFallback :: event -> currentState -> Either String State

instance {-# OVERLAPPABLE #-} NextStateWithFallbackInstance s e where
  nextWFallback _ _ = Left "Invalid Transition"

instance NextStateWithFallbackInstance Init GotInt where
  nextWFallback (GotInt i) Init = Right $ StateA (A i)
instance NextStateWithFallbackInstance Init GotString where
  nextWFallback (GotString s) Init = Right $ StateB (B s)

instance NextStateWithFallbackInstance A GotString where
  nextWFallback (GotString s) (A i) = Right $ StateC (C i s)

instance NextStateWithFallbackInstance B GotInt where
  nextWFallback (GotInt i) (B s) = Right $ StateC (C i s)

instance NextStateWithFallbackInstance State Event where
  -- We could also split this dispatch table into several, by adding an instance
  -- for each State that would handle the Event sum type.
  nextWFallback (EventGotInt e)    (StateInit s) = nextWFallback e s
  nextWFallback (EventGotString e) (StateInit s) = nextWFallback e s
  nextWFallback (EventGotInt e)    (StateA s)    = nextWFallback e s
  nextWFallback (EventGotString e) (StateA s)    = nextWFallback e s
  nextWFallback (EventGotInt e)    (StateB s)    = nextWFallback e s
  nextWFallback (EventGotString e) (StateB s)    = nextWFallback e s
  nextWFallback (EventGotInt e)    (StateC s)    = nextWFallback e s
  nextWFallback (EventGotString e) (StateC s)    = nextWFallback e s

foldEventsMPTCWFallback :: State -> [Event] -> Either String State
foldEventsMPTCWFallback = foldlM (flip nextWFallback)

-- Third approach:
-- One single-param type class per event
-- We still have to wire up this "dispatch function" by hand
-- The type classes ensure we cannot add in a bad transition.
-- We could, however, forget to use a valid transition, in which case it would
-- fall through to the error case.
-- One advantage here is that we could get rid of the "GotInt" and "GotString"
-- types if we wanted, unpacking their content directly into the Event sum type.
class HandleGotInt currentState where
  gotInt :: GotInt -> currentState -> State

class HandleGotString currentState where
  gotString :: GotString -> currentState -> State

instance HandleGotInt Init where
  gotInt (GotInt i) Init = StateA (A i)
instance HandleGotString Init where
  gotString (GotString s) Init = StateB (B s)

instance HandleGotString A where
  gotString (GotString s) (A i) = StateC (C i s)

instance HandleGotInt B where
  gotInt (GotInt i) (B s) = StateC (C i s)

class HandleEvent currentState where
  handleEvent :: Event -> currentState -> Either String State

instance HandleEvent Init where
  handleEvent (EventGotInt e)    = Right . gotInt e
  handleEvent (EventGotString e) = Right . gotString e

instance HandleEvent A where
  handleEvent (EventGotString e) s = Right $ gotString e s
  handleEvent _                  _ = Left "Invalid Transition"

instance HandleEvent B where
  handleEvent (EventGotInt e) s = Right $ gotInt e s
  handleEvent _               _ = Left "Invalid Transition"

instance HandleEvent C where
  handleEvent _ _ = Left "Invalid Transition"

instance HandleEvent State where
  handleEvent e (StateInit s) = handleEvent e s
  handleEvent e (StateA s) = handleEvent e s
  handleEvent e (StateB s) = handleEvent e s
  handleEvent e (StateC s) = handleEvent e s

foldEventsOneTypeclassPerEvent :: State -> [Event] -> Either String State
foldEventsOneTypeclassPerEvent = foldlM (flip handleEvent)

-- | Fourth approach:
-- Multi-param type class with a third parameter to point the result state
-- In this case, we get pretty strong guarantees if we know the types of all
-- states and events at compile time (check the examples at the end).
--
-- Notice that in this approach the "content" of the state cannot be used to
-- choose the final state; the type of the current state and event fully determine
-- the type of the next state. This is expressed by the functional dependency on
-- the Transition class. This functional dependency forbids us to create two
-- instances for the same pair of currentState and event types with different
-- resulting states. On the other hand, it allows the compiler to infer the type
-- of a call to transitionTo.
--
-- The sum-type instances allow us to go back to the normal approach of folding
-- a list of Events on top of some State.
-- The AsState typeclass isn't really necessary, but it allows the sum-type instances
-- to avoid referring directly to the State constructors, so if we change the result
-- type of some transition, we do not have to change the sum-type instance.
class Transition currentState event nextState | currentState event -> nextState where
  transitionTo :: event -> currentState -> nextState

-- Specific instances
instance Transition Init GotInt A where
  transitionTo (GotInt i) Init = A i
instance Transition Init GotString B where
  transitionTo (GotString s) Init = B s

instance Transition A GotString C where
  transitionTo (GotString s) (A i) = C i s

instance Transition B GotInt C where
  transitionTo (GotInt i) (B s) = C i s

-- Cool thing we can do at compile time
(~:>) :: Transition state event next => state -> event -> next
state ~:> event = transitionTo event state
infixl 7 ~:>

(>>:=) :: (Monad m, Transition state event next) => m state -> (state -> m event) -> m next
(>>:=) mstate fevent = do
  state <- mstate
  event <- fevent state
  pure (transitionTo event state)
infixl 2 >>:=

(<*:>) :: (Monad m, Transition state event next) => m state -> m event -> m next
(<*:>) mstate mevent = flip transitionTo <$> mstate <*> mevent
infixl 2 <*:>

x :: C
x = Init ~:> GotInt 42 ~:> GotString "The Answer"

y :: IO C
y = pure Init
    <*:> GotInt <$> readLn
    <*:> GotString <$> getLine

z :: IO C
z = pure Init
    <*:> GotInt <$> readLn
    >>:= \case (A i) | i == 42   -> pure $ GotString "The Answer"
                     | otherwise -> GotString <$> getLine

-- Sum-type instances
instance Transition Init Event (Either String State) where
  transitionTo (EventGotInt e) s    = Right . asState $ transitionTo e s
  transitionTo (EventGotString e) s = Right . asState $ transitionTo e s

instance Transition A Event (Either String State) where
  transitionTo (EventGotString e) s = Right . asState $ transitionTo e s
  transitionTo _ _ = Left "Invalid transition"

instance Transition B Event (Either String State) where
  transitionTo (EventGotInt e) s = Right . asState $ transitionTo e s
  transitionTo _ _ = Left "Invalid transition"

instance Transition C Event (Either String State) where
  transitionTo _ _ = Left "Invalid transition"

instance Transition State Event (Either String State) where
  transitionTo e (StateInit s) = transitionTo e s
  transitionTo e (StateA s)    = transitionTo e s
  transitionTo e (StateB s)    = transitionTo e s
  transitionTo e (StateC s)    = transitionTo e s

-- Helper class to wrap states
class AsState s where
  asState :: s -> State

instance AsState Init where
  asState = StateInit
instance AsState A where
  asState = StateA
instance AsState B where
  asState = StateB
instance AsState C where
  asState = StateC

foldEventsMPTC2 :: State -> [Event] -> Either String State
foldEventsMPTC2 = foldlM (flip transitionTo)