statechart2.hs

StateChart.hs

{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE DerivingVia                #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TupleSections              #-}
{-# LANGUAGE TypeApplications           #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module StateChart where

import           Data.Map (Map)
import qualified Data.Map as M


------------------------------------------------------------------------------
-- | Before we start; a few instances that should exist in 'base', but don't.
instance (Bounded b, Enum a, Enum b) => Enum (a, b) where
  fromEnum (a, b) = (fromEnum (maxBound @b) + 1) * fromEnum a + fromEnum b
  toEnum n =
    let bound = fromEnum (maxBound @b) + 1
        b = n `rem` bound
        a = n `div` bound
     in (toEnum a, toEnum b)

instance (Bounded a, Bounded b) => Bounded (Either a b) where
  minBound = Left minBound
  maxBound = Right maxBound

instance (Bounded a, Enum a, Enum b) => Enum (Either a b) where
  toEnum i =
    let bound = fromEnum (maxBound @a) + 1
     in case i < bound of
          True -> Left $ toEnum i
          False -> Right $ toEnum $ i - bound
  fromEnum (Left a) = fromEnum a
  fromEnum (Right b) = fromEnum b + fromEnum (maxBound @a) + 1



------------------------------------------------------------------------------
-- Nodes
------------------------------------------------------------------------------


------------------------------------------------------------------------------
-- | The things we can do in a given state.
--
-- The type parameters are, in order:
--     m: the monad in which we can invoke actions
--     e: the events we allow
--     s: the possible states
--     a: the eventual value we accept
data Node m e s a where
  -- | Either accept the state, producing a final result
  Accept :: a -> Node m e s a
  -- | Or transition to another state, based on some event 'e'
  Transition :: (e -> s) -> Node m e s a
  -- | Or run a monadic computation, and branch to a new state based on its
  -- result
  Invoke  :: (Bounded x, Enum x, Show x) => String -> m x -> (x -> s) -> Node m e s a

instance Functor (Node m e s) where
  fmap fab (Accept a) = Accept (fab a)
  fmap _ (Transition fes) = Transition fes
  fmap _ (Invoke s mx fxs) = Invoke s mx fxs


------------------------------------------------------------------------------
-- | Given a 'Node' and an event, evaluate the node, returning either a new state
-- or a final value.
runNode :: Applicative m => Node m e s a -> e -> m (Either s a)
runNode n e = case n of
   Accept a -> pure $ Right a
   Transition fenmesa -> pure $ Left $ fenmesa e
   -- TODO(sandy): There's a bug here; this silently eats the event. Not hard
   -- to fix, but only noticed it when writing documentation.
   Invoke _ m f -> fmap (Left . f) m

------------------------------------------------------------------------------
-- | Given a 'Node', determine if we accepted with a final result.
finalize :: Node m e s a -> Maybe a
finalize (Accept a) = Just a
finalize (Transition _) = Nothing
finalize (Invoke _ _ _) = Nothing


------------------------------------------------------------------------------
-- | Zip two 'Node's together. The result is a 'Node' that has both states, and
-- can accept events for either 'Node'.
--
-- This is a generalization of the 'Applicative' method @'liftA2' (,)@.
tensor
    :: s1
    -> s2
    -> Node m e1 s1 a
    -> Node m e2 s2 b
    -> Node m (Either e1 e2) (s1, s2) (a, b)
tensor _  _  (Accept a) (Accept b)
  = Accept (a, b)
tensor ak bk (Accept _) (Transition fe2s2)
  = Transition $ either (const (ak, bk)) ((ak, ) . fe2s2)
tensor ak _  (Accept _) (Invoke s mx fxs2)
  = Invoke ("2 " <> s) mx $ (ak, ) . fxs2
tensor ak bk (Transition fe1s1) (Accept _)
  = Transition $ either ((, bk) . fe1s1) (const (ak, bk))
tensor ak bk (Transition fe1s1) (Transition fe2s2)
  = Transition $ either ((, bk) . fe1s1) ((ak, ) . fe2s2)
tensor ak _  (Transition _) (Invoke s mx fxs2)
  = Invoke ("2 " <> s) mx $ (ak, ) . fxs2
tensor _  bk (Invoke s mx fxs1) _
  = Invoke ("1 " <> s) mx $ (, bk) . fxs1



------------------------------------------------------------------------------
-- State Charts
------------------------------------------------------------------------------

------------------------------------------------------------------------------
-- | A 'StateChart' is a mapping from states to 'Node's.
newtype StateChart m e s a = StateChart
  { unStateChart :: Map s (Node m e s a)
  }
  deriving newtype (Semigroup, Monoid)
  deriving stock Functor


------------------------------------------------------------------------------
-- | Find the corresponding 'Node' to a state in a 'StateChart'.
lookupS :: Ord s => StateChart m e s a -> s -> Maybe (Node m e s a)
lookupS (StateChart sc) s = M.lookup s sc


------------------------------------------------------------------------------
-- | Run a 'StateChart' given a series of events and a starting state. Performs
-- all invoked actions, and then returns either the resulting state or the
-- final accepted value.
run :: (Ord s, Monad m) => StateChart m e s a -> [e] -> s -> m (Either s a)
run sc es0 s = go es0 $ lookupS sc s
  where
    go [] (Just x) = pure $ note s $ finalize x
    go (e : es) (Just nmess) = either (run sc es) (pure . Right) =<< runNode nmess e
    go _ Nothing = pure $ Left s


------------------------------------------------------------------------------
-- | Lift 'tensor' over two 'StateChart's.
tensorSC
    :: (Ord s1, Ord s2)
    => StateChart m e1 s1 a
    -> StateChart m e2 s2 b
    -> StateChart m (Either e1 e2) (s1, s2) (a, b)
tensorSC (StateChart sca) (StateChart scb) = StateChart $ M.fromList $ do
  (ak, av) <- M.toList sca
  (bk, bv) <- M.toList scb
  pure $ ((ak, bk), tensor ak bk av bv)


------------------------------------------------------------------------------
-- | Raise a 'Maybe' to an 'Either'.
note :: s -> Maybe a -> Either s a
note s Nothing = Left s
note _ (Just a) = Right a


------------------------------------------------------------------------------
-- Helper functions for defining state charts
------------------------------------------------------------------------------

transition :: s -> (e -> s) -> StateChart m e s a
transition s = StateChart . M.singleton s . Transition


terminal :: s -> a -> StateChart m e s a
terminal s = StateChart . M.singleton s . Accept


invoke
    :: (Enum x, Bounded x, Show x)
    => s -> String -> m x -> (x -> s) -> StateChart m e s a
invoke s lbl m = StateChart . M.singleton . $ Invoke lbl m



------------------------------------------------------------------------------
-- Example state charts
------------------------------------------------------------------------------

------------------------------------------------------------------------------
-- | A 'StateChart' that immediately accepts the first 'Bool' event it
-- receives.
boolean :: StateChart m Bool (Maybe Bool) Bool
boolean = mconcat
  [ transition Nothing Just
  , terminal (Just True) True
  , terminal (Just False) False
  ]


------------------------------------------------------------------------------
-- | A 'StateChart' that computes '(&&)' by running two 'boolean's in parallel.
andSC :: StateChart m (Either Bool Bool) (Maybe Bool, Maybe Bool) Bool
andSC = fmap (uncurry (&&)) $ tensorSC boolean boolean



------------------------------------------------------------------------------
-- Machinery for inspecting 'StateChart's' via graphviz.
------------------------------------------------------------------------------

arr :: (Show src, Show dst, Show e) => src -> dst -> Maybe e -> String
arr s s' e = arr' (show $ show s) (show $ show s') (fmap (show . show) e)


arr' :: String -> String -> Maybe String -> String
arr' s s' (Just e) = mconcat
  [ s , " -> " , s' , " [label=" , e , "];" ]
arr' s s' Nothing = mconcat
  [ s , " -> " , s' , " [style=dotted];" ]



inspect :: (Show a, Bounded e, Enum e, Show e, Show s) => StateChart m e s a -> String
inspect sc
  = unlines
  . flip mappend ["}"]
  . ("digraph x {" :)
  . foldMap (\(s, x) -> inspect' (show s) x)
  . M.toList
  $ unStateChart sc


node :: Show s => s -> s -> Maybe String -> String
node nm lbl Nothing = mconcat [ show nm, "[label=", show lbl, "];" ]
node nm lbl (Just l_c) = mconcat [ show nm, " [label=" <> show lbl <> ",shape=", l_c, "];" ]


type Label = String


inspect'
    :: forall s e a m
     . (Show a, Bounded e, Enum e, Show e, Show s)
    => Label
    -> Node m e s a
    -> [String]
inspect' lbl (Accept a) = [arr lbl a $ Nothing @(), node a a $ Just "box"]
inspect' lbl (Transition fenmesa) = do
  e <- [minBound .. maxBound]
  let t = fenmesa e
  pure $ arr lbl (show t) (Just e)
inspect' lbl (Invoke s _ fxnmesa) =
  node (show (lbl <> s)) s (Just "house") : arr lbl (lbl <> s) (Nothing @()) : do
    x <- [minBound .. maxBound]
    let t = fxnmesa x
    pure $ arr (lbl <> s) (show t) $ Just x

inspect $ andSC @IO: