Type-safe extensible state machine https://stackoverflow.com/questions/50579397/extensible-state-machines-in-haskell/50588021#50588021

SM.hs

{-# LANGUAGE
    EmptyCase,

    DataKinds,
    PolyKinds,
    KindSignatures,

    GADTs,
    FlexibleContexts,
    FlexibleInstances,
    MultiParamTypeClasses,
    TypeFamilies,
    UndecidableInstances,

    AllowAmbiguousTypes,
    TypeApplications,
    ScopedTypeVariables,

    StandaloneDeriving
  #-}

module SM where

import Data.Kind
import GHC.TypeLits (Symbol)

-- * Extensible sums

data S (u :: [(kk, Type)]) where
  Here :: forall k a u. a -> S ('(k, a) ': u)
  There :: forall u x. S u -> S (x ': u)

deriving instance (Show a, Show (S u)) => Show (S ('(k, a) ': u))
instance Show (S '[]) where
  showsPrec _ x = case x of {}

-- Injection into a sum.
-- v is a list containing the pair (k, a)
class Inj1 (k :: kk) (a :: Type) (v :: [(kk, Type)]) where
  inj1 :: a -> S v

instance {-# OVERLAPPING #-} (a ~ a') => Inj1 k a ('(k, a') ': _v) where
  inj1 = Here

instance Inj1 k a v => Inj1 k a ('(_h, _b) ': v) where
  inj1 = There . inj1 @_ @k

-- Make up a new key if none of the other matches
instance {-# INCOHERENT #-} (v ~ ('(k, a) ': w)) => Inj1 k a v where
  inj1 = Here

-- Injection from a sum u into a bigger sum v
class Inj (u :: [(kk, Type)]) (v :: [(kk, Type)]) where
  inj :: S u -> S v

instance Inj '[] v where
  inj x = case x of {}

instance (Inj1 k a v, Inj u v) => Inj ('(k, a) ': u) v where
  inj (Here a) = inj1 @_ @k a
  inj (There b) = inj b

-- * State machines

-- @k@ is a state, with memory @Contents k@ (i.e., runtime content)
-- and output states @Outputs k@.
class State (k :: kk) where
  type Contents k :: Type
  type Outputs k :: [(kk, Type)]
  transition :: Contents k -> S (Outputs k)

-- Construct a state machine with initial state @i@.
class StateMachine (i :: kk) (u :: [(kk, Type)]) where
  sm :: S u -> S u

-- ** Internals

instance (u ~ ('(i, a) ': v), SMBuilder u u) => StateMachine i u where
  sm = smb

-- Internal auxiliary definition
class SMBuilder (u :: [(kk, Type)]) (v :: [(kk, Type)]) where
  smb :: S v -> S u

instance (State k, a ~ Contents k, Inj (Outputs k) u, SMBuilder u v) => SMBuilder u ('(k, a) ': v) where
  smb (Here a) = inj (transition @_ @k a)
  smb (There x) = smb x

instance {-# INCOHERENT #-} (v ~ '[]) => SMBuilder u v where
  smb x = case x of {}

-- * Example

-- Declare a namespace of state names for our example application
data Named (s :: Symbol)

-- ** Module A (defining state A)

instance State (Named "A") where
  -- No memory in the state.
  type Contents (Named "A") = ()
  -- There is only a transition to "B"
  type Outputs (Named "A") = '[ '(Named "B", Int) ]

  transition () = inj1 @_ @(Named "B") 10

-- ** Module B (defining state B)

instance State (Named "B") where
  type Contents (Named "B") = Int
  type Outputs (Named "B") = '[ '(Named "A", ()), '(Named "B", Int)]
  transition i
    | i < 0 = inj1 @_ @(Named "A") ()
    | otherwise = inj1 @_ @(Named "B") (i-1)

-- ** Module Main

-- "A" is the initial state
type Initial = Named "A"

-- Initial value in state "A"
initial :: Inj1 Initial () u => S u
initial = inj1 @_ @Initial ()

-- Construct state machine
transition' = sm @_ @(Named "A")

main :: IO ()
main = do
  let steps = 14
  (mapM_ print . take (steps+1) . iterate transition') initial

{-
  Here ()
  There (Here 10)
  There (Here 9)
  There (Here 8)
  There (Here 7)
  There (Here 6)
  There (Here 5)
  There (Here 4)
  There (Here 3)
  There (Here 2)
  There (Here 1)
  There (Here 0)
  There (Here (-1))
  Here ()
  There (Here 10)
-}