State machines as natural transformations between (polynomial) functors, a la David Spivak

StateMachine.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TypeOperators #-}
module StateMachine where

import Control.Comonad (extract)
import Control.Comonad.Cofree
import Data.Bifunctor (first, second)
import Data.Functor.Day
import GHC.Generics

infixr 0 ~>
type p ~> q = forall y. p y -> q y

class HBifunctor f where
  hbimap :: (Functor p, Functor q) => (p ~> p') -> (q ~> q') -> f p q ~> f p' q'

instance HBifunctor (:.:) where
  hbimap fp fq = Comp1 . fp . fmap fq . unComp1

instance HBifunctor Day where
  hbimap fp fq (Day p q f) = Day (fp p) (fq q) f

instance HBifunctor (:*:) where
  hbimap fp fq (p :*: q) = (fp p :*: fq q)


-- A monomial a*y^b
data Mono a b y = Mono a (b -> y) deriving Functor

splitMonoComp :: Mono (a, c) (b, d) ~> Mono a b :.: Mono c d
splitMonoComp (Mono (a, c) bdy) = Comp1 (Mono a (\b -> Mono c (\d -> bdy (b, d))))

splitMonoDay :: Mono (a, c) (b, d) ~> Mono a b `Day` Mono c d
splitMonoDay (Mono (a, c) f) = Day (Mono a id) (Mono c id) (curry f)

splitMonoProd :: Mono (a, b) (a, b) ~> Mono a a :*: Mono b b
splitMonoProd (Mono (a, b) f) = Mono a (\a' -> f (a', b)) :*: Mono b (\b' -> f (a, b'))


-- A state machine with internal state
data SM p = forall s. SM { transition :: Mono s s ~> p, initialState :: s }

mapSM :: p ~> q -> SM p -> SM q
mapSM pq (SM f s) = SM (pq . f) s

runSM :: SM p -> p (SM p)
runSM = runMapSM id

runMapSM :: (SM p -> a) -> SM p -> p a
runMapSM g (SM f s) = f (Mono s (g . SM f))

runMultipleSM :: SM p -> Cofree p (SM p)
runMultipleSM sm = sm :< runMapSM runMultipleSM sm

basicSM :: Functor p => (s -> p s) -> s -> SM p
basicSM f = SM (\(Mono s g) -> g <$> f s)

combineSM
  :: (HBifunctor comb, Functor p, Functor q)
  => (forall a b. Mono (a, b) (a, b) ~> Mono a a `comb` Mono b b)
  -> SM p -> SM q -> SM (p `comb` q)
combineSM splitMono (SM fp sp) (SM fq sq) = SM (hbimap fp fq . splitMono) (sp, sq)

seqSM :: (Functor p, Functor q) => SM p -> SM q -> SM (p :.: q)
seqSM = combineSM splitMonoComp

parSM :: (Functor p, Functor q) => SM p -> SM q -> SM (p `Day` q)
parSM = combineSM splitMonoDay

altSM :: (Functor p, Functor q) => SM p -> SM q -> SM (p :*: q)
altSM = combineSM splitMonoProd


newtype Mealy i o y = Mealy { unMealy :: i -> (o, y) }
  deriving Functor

combineCompMealy :: Mealy a b :.: Mealy b c ~> Mealy a c
combineCompMealy (Comp1 (Mealy abm)) = Mealy $ \a -> case abm a of (b, Mealy bcy) -> bcy b

combineDayMealy :: Mealy a b `Day` Mealy c d ~> Mealy (a, c) (b, d)
combineDayMealy (Day (Mealy aby) (Mealy cdy) f) = Mealy $
  \(a, c) -> case (aby a, cdy c) of
    ((b, ya), (d, yc)) -> ((b, d), f ya yc)

combineProdMealy :: Mealy a b :*: Mealy c d ~> Mealy (Either a c) (Either b d)
combineProdMealy (Mealy aby :*: Mealy cdy) = Mealy $ either (first Left . aby) (first Right . cdy)


basicMeally :: (s -> i -> (o, s)) -> s -> SM (Mealy i o)
basicMeally f = basicSM (Mealy . f)

runMealy :: SM (Mealy i o) -> i -> (o, SM (Mealy i o))
runMealy = unMealy . runSM

runMultipleMealy :: (Monoid o, Foldable f) => SM (Mealy i o) -> f i -> (o, SM (Mealy i o))
runMultipleMealy c = second extract . foldl (\(o, _ :< Mealy ioc) -> first (o <>) . ioc) (mempty, runMultipleSM c)

seqMealy :: SM (Mealy a b) -> SM (Mealy b c) -> SM (Mealy a c)
seqMealy s0 s1 = mapSM combineCompMealy $ seqSM s0 s1

parMealy :: SM (Mealy a b) -> SM (Mealy c d) -> SM (Mealy (a, c) (b, d))
parMealy s0 s1 = mapSM combineDayMealy $ parSM s0 s1

altMealy :: SM (Mealy a b) -> SM (Mealy c d) -> SM (Mealy (Either a c) (Either b d))
altMealy s0 s1 = mapSM combineProdMealy $ altSM s0 s1