Demonstration of a very-type-safe finite deterministic state machine using DataKinds, ConstraintKinds and TypeFamilies

StateMachine2.hs

-- Demonstration of a finite deterministic state machine
-- The state is tracked using a single data type, which is parametrised
-- over the type of messages which are allowed to update the state at any
-- point in time.
-- This should be more clear when reading the code :-)

-- Kudos to elliott, merlijn, mm_freak and others in #haskell for all help
-- and pointers

{-# LANGUAGE DataKinds,
             TypeOperators,
             KindSignatures,
             GADTs,
             StandaloneDeriving,
             TypeFamilies,
             ConstraintKinds,
             UndecidableInstances,
             PolyKinds,
             Rank2Types #-}

module StateMachine2 where

-- Part 1: The heavy plumbing (aka magic) to get things to work later on

-- For 'Constraint'
import GHC.Exts

-- A hack to turn the type-level 'elem' defined below into a GHC
-- 'Constraint', which allows for nicer syntax in our 'user code'
type family Elem (x :: k) (xs :: [k]) :: Constraint
type instance Elem x xs = (Elem' x xs ~ ValidMessageForState)

data AcceptableMessage = ValidMessageForState
                       | InvalidMessageForState

-- A type family to encode 'elem' on type-level lists
type family Elem' (x :: k) (xs :: [k]) :: AcceptableMessage
type instance where
    Elem' x '[] = InvalidMessageForState
    Elem' x (x ': xs) = ValidMessageForState
    Elem' x (y ': xs) = Elem' x xs

-- End of part 1. Rather short, huh?

-- Part 2: Implementation of our state-machine
-- Types of messages which can be pushed through the state
data Message1 = Message1 Int
data Message2 = Message2 Int

-- The state data type. Notice it's parametrised over a list of types.
-- For this demo, the state only contains a single Int value
data State :: [*] -> * where
    State :: Int -> State m

deriving instance Show (State m)

type Handler m n = Elem m ms => State ms -> m -> State n

-- Handler for 'Message1' messages. The constraint defines only states
-- which are expected to accept a Message1, can be passed through the
-- function.
-- After completion, this function returns a state which can be updated
-- using both a Message1 or Message2 message.
handleMessage1 :: Handler Message1 '[Message1, Message2]
handleMessage1 (State i) (Message1 m) = State (i + m)

-- Similar to handleMessage1, handleMessage2 requires a state which can
-- accept an update through a Message2 value. Unlike handleMessage1, it
-- returns a state which can only be updated through a Message1 message.
handleMessage2 :: Handler Message2 '[Message1]
handleMessage2 (State i) (Message2 m) = State (i - m)

-- Some initial state. Notice due to its type, this can only be passed to
-- 'handleMessage1', not 'handleMessage2'.
initialState :: State '[Message1]
initialState = State 0

-- Testing code!
main :: IO ()
main = do
    let s0 = initialState
        s1 = handleMessage1 s0 (Message1 10)

        -- Due to the type of s1 (State '[Message1, Message2]), we can pass
        -- it to both handleMessage1 and handleMessage2
        s2 = handleMessage2 s1 (Message2 4)
        _s2' = handleMessage1 s1 (Message1 5)

        -- Due to the type of s2 (State '[Message1]), it can be passed to
        -- handleMessage1 but not handleMessage2
        s3 = handleMessage1 s2 (Message1 5)

        -- This doesn't type-check (as desired):
        -- s3' = handleMessage2 s2 (Message2 10)
        --
        -- The error:
        --
        --   Couldn't match type 'InvalidMessageForState
        --           with 'ValidMessageForState
        --   In the expression: handleMessage2 s2 (Message2 10)
        --   In an equation for s3': s3' = handleMessage2 s2 (Message2 10)
        --   In the expression:
        --     do { let s0 = initialState
        --              s1 = handleMessage1 s0 (Message1 10)
        --              ....;
        --          print s3 }
        --
        -- Rationale: s2 is the result of handleMessage2, and such state should
        -- only be used with Message1

    print s3 -- State 11