TH.hs

{-# LANGUAGE TemplateHaskell #-}
-- | The @scxml@ quasiquoter. Used at the top level of a module:
--
-- @
-- [scxml| <scxml initial="Draft"> ... </scxml> |]
-- @
--
-- The generated names are fixed, so a module holds one chart:
--
-- * @data FsmState@: one constructor per child of @<scxml>@, named exactly as
--   the state id. A compound state becomes a constructor carrying a sum type
--   of the same name as the state; a parallel state becomes a constructor
--   with one field per compound region; atomic and final states are nullary.
-- * @data FsmEvent@: one constructor per event name, verbatim, carrying the
--   payload types written after the name in the @event@ attribute, plus a
--   nullary @DoneX@ for SCXML's automatic @done.state.X@ completion events.
-- * @data FsmTransition@: one step of the chart, a record of
--   @transitionFrom :: FsmState@, @transitionEvent :: FsmEvent@ and
--   @transitionTo :: FsmState@. What an exit callback receives.
-- * @data FsmEntry = FsmStart FsmState | FsmStep FsmTransition@: what an
--   entry callback receives. @FsmStart@ carries the initial state while
--   @initiateStateMachine@ enters it, which leaves no state and has no event.
-- * @fsmChart :: Def FsmState FsmEvent@, for "Scxml.Statechart.Run".
-- * @serializeStateMachine :: FsmState -> [Text]@ and
--   @deserializeStateMachine :: [Text] -> Maybe FsmState@, which store a state
--   as the set of active state ids and read it back, rejecting anything that
--   is not a configuration of this chart.
-- * @initiateStateMachine@ and @notifyStateMachine@, which run the chart
--   calling the callbacks named in @<script>@ elements. Callbacks are looked
--   up by name in the module containing the quasiquote (they may be defined
--   after it) and must share one monad, which the type checker enforces:
--
-- @
-- initiateStateMachine :: m FsmState
-- notifyStateMachine   :: FsmState -> FsmEvent -> m FsmState
--
-- entryCallback :: FsmEntry -> m (Maybe FsmEvent)
-- exitCallback  :: FsmTransition -> m ()
-- @
--
-- Every generated state type derives @Show@, @Read@, @Eq@ and @Ord@, and the
-- event type also derives @Enum@ and @Bounded@ as long as no event carries a
-- payload: @Ord@ would demand an instance of every payload type, and the
-- other two need every constructor nullary. @FsmTransition@ and @FsmEntry@
-- derive @Show@, @Read@ and @Eq@, which every event type has. For storing a state outside
-- Haskell, prefer the generated @serializeStateMachine@ over @Show@.
module Scxml.Statechart.TH (scxml) where

import Control.Monad (forM, unless)
import Data.Char (isAlphaNum, isLower, isUpper)
import Data.List (nub)
import qualified Data.Map.Strict as Map
import qualified Data.List.NonEmpty as NE
import qualified Data.Set as Set
import Data.Text (Text)
import qualified Data.Text as T
import Language.Haskell.TH
import Language.Haskell.TH.Quote
import Language.Haskell.TH.Syntax (lift)

import Scxml.Statechart.Def
import qualified Scxml.Statechart.Interpret as I
import Scxml.Statechart.Model
import Scxml.Statechart.Parse
import qualified Scxml.Statechart.Run as Run

-- | Turns SCXML into a statechart. Usable only at the top level of a module,
-- where it declares the state and event types and the functions that run the
-- chart. A light switch, whose two states each announce themselves on entry:
--
-- > {-# LANGUAGE QuasiQuotes #-}
-- > module LightSwitch where
-- >
-- > import Scxml.Statechart (scxml)
-- >
-- > [scxml|
-- > <scxml initial="Off">
-- >   <state id="Off">
-- >     <onentry><script>report</script></onentry>
-- >     <transition event="Flip" target="On"/>
-- >   </state>
-- >   <state id="On">
-- >     <onentry><script>report</script></onentry>
-- >     <transition event="Flip" target="Off"/>
-- >   </state>
-- > </scxml>
-- > |]
-- >
-- > report :: FsmEntry -> IO (Maybe FsmEvent)
-- > report (FsmStart st) = print st >> pure Nothing
-- > report (FsmStep step) = print (transitionTo step) >> pure Nothing
--
-- State ids and event names are used verbatim, so that declares @FsmState@
-- with constructors @Off@ and @On@, @FsmEvent@ with @Flip@, and:
--
-- > initiateStateMachine :: IO FsmState
-- > notifyStateMachine   :: FsmState -> FsmEvent -> IO FsmState
--
-- Callbacks are ordinary functions in the same module, written after the
-- quasiquote, and are found by the name in @\<script\>@. An exit callback
-- receives an @FsmTransition@: the state left, the event that caused the
-- step, and the state the step settles in. An entry callback receives an
-- @FsmEntry@, which is @FsmStep@ with that same record, or @FsmStart@ with
-- the initial state while @initiateStateMachine@ enters it:
--
-- > ghci> off <- initiateStateMachine
-- > Off
-- > ghci> on <- notifyStateMachine off Flip
-- > On
--
-- The callbacks all share one monad, which the type checker enforces. @IO@ is
-- enough to print; when a chart has to carry something along, put the
-- callbacks in a @StateT@ over it and the generated functions follow suit:
--
-- > record :: FsmEntry -> StateT [FsmState] IO (Maybe FsmEvent)
-- > record entry = modify (entered :) >> pure Nothing
-- >   where entered = case entry of FsmStart st -> st; FsmStep step -> transitionTo step
-- >
-- > -- initiateStateMachine :: StateT [FsmState] IO FsmState
-- > -- notifyStateMachine   :: FsmState -> FsmEvent -> StateT [FsmState] IO FsmState
--
-- Returning @Just event@ instead of @Nothing@ raises that event, which is how
-- a callback decides where the chart goes next.
--
-- An event may carry data, by writing the Haskell types its constructor holds
-- after its name. @\<transition event="Order Item Int" target="Checking"/\>@
-- declares @Order Item Int@, and the callbacks that transition runs receive
-- the value the caller passed in, payload and all, as @transitionEvent@. A payload
-- type is one type constructor, optionally qualified (@Int@,
-- @Order.LineItem@); like a callback name it is resolved after the
-- quasiquote, so it may be defined below it. Transitions are still selected
-- by event name alone, so a payload never decides where the chart goes; that
-- stays with the events a callback raises.
--
-- The generated names are fixed, so a module holds one chart.
scxml :: QuasiQuoter
scxml =
  QuasiQuoter
    { quoteExp = const unsupported
    , quotePat = const unsupported
    , quoteType = const unsupported
    , quoteDec = generate
    }
  where
    unsupported = fail "scxml: only usable as a top-level declaration, e.g. [scxml| <scxml ...> |]"

-- | A compound-like node (the root or a @<state>@ with children) becomes its
-- own data type with helpers converting to and from configurations.
data Group = Group
  { gType     :: Name
  , gChildren :: [Node]
  , gTo       :: Name
  , gFrom     :: Name
  }

generate :: String -> Q [Dec]
generate src = do
  ch <- orFail (parseScxml src)
  -- Everything here is a walk of the tree: a node carries its own children,
  -- so nothing needs looking up by id.
  let states = chartStates ch
  let stateT = mkName "FsmState"
      eventT = mkName "FsmEvent"
      transitionT = mkName "FsmTransition"
      entryT = mkName "FsmEntry"
      startCon = mkName "FsmStart"
      stepCon = mkName "FsmStep"
      fromField = mkName "transitionFrom"
      eventField = mkName "transitionEvent"
      toField = mkName "transitionTo"
      defName = mkName "fsmChart"
      startName = mkName "initiateStateMachine"
      stepName = mkName "notifyStateMachine"
      toIdsName = mkName "serializeStateMachine"
      fromIdsName = mkName "deserializeStateMachine"
      -- A compound state's type has the same name as its constructor; Haskell
      -- keeps types and constructors in separate namespaces.
      nameFor sid = mkName (T.unpack sid)
      compounds = [n | n <- states, Compound _ <- [nodeKind n]]

  -- Callback names
  let entryActions = nub (concatMap nodeOnEntry states)
      exitActions = nub (concatMap nodeOnExit states)
  mapM_ (orFail . checkVarName) (nub (entryActions ++ exitActions))

  -- Events: those named in transitions, in document order, then done events of
  -- states that can complete but that no transition mentions.
  let referenced = chartEvents ch
      completing = [n | n <- states, completes n]
      doneEvents = [Event (I.doneEventName nodeId) [] | Node {nodeId} <- completing]
      events = referenced ++ filter ((`notElem` map eventName referenced) . eventName) doneEvents
      -- (constructor, payload type names, what it came from, wire name)
      eventInfo Event {eventName = name, eventFields} =
        case T.stripPrefix (T.pack "done.state.") name of
          Just sid -> (mkName ("Done" ++ T.unpack sid), [], "completion event " ++ show (T.unpack name), name)
          Nothing -> (mkName (T.unpack name), eventFields, "event " ++ show (T.unpack name), name)
      eventInfos = map eventInfo events
      eventCons = [(c, origin) | (c, _, origin, _) <- eventInfos]
      -- Enum and Bounded need every constructor nullary, and Ord would make
      -- the whole chart fail to compile over a payload type that has no
      -- instance, so an event carrying data costs all three.
      eventDerivs
        | null eventInfos = [''Show, ''Read, ''Eq, ''Ord]
        | any (\(_, fs, _, _) -> not (null fs)) eventInfos = [''Show, ''Read, ''Eq]
        | otherwise = [''Show, ''Read, ''Eq, ''Ord, ''Enum, ''Bounded]

  groups <- forM (Nothing : map Just compounds) $ \g -> do
    let ty = maybe stateT (nameFor . nodeId) g
    to <- newName ("toCfg_" ++ nameBase ty)
    from <- newName ("fromCfg_" ++ nameBase ty)
    let children = maybe (NE.toList (chartRoot ch)) nodeChildren g
    pure (g, Group ty children to from)
  let groupMap = Map.fromList [(nodeId, grp) | (Just Node {nodeId}, grp) <- groups]
  Group {gTo = rootTo, gFrom = rootFrom} <- case groups of
    (_, g) : _ -> pure g
    [] -> fail "scxml: internal error, no root group"
  let groupOf sid = case Map.lookup sid groupMap of
        Just grp -> pure grp
        Nothing -> fail ("scxml: internal error, no group for " ++ T.unpack sid)

  -- Every generated constructor and type, with its origin, so clashes give a
  -- readable error instead of "Multiple declarations".
  let stateCons =
        [ (nameFor nodeId, "state " ++ show (T.unpack nodeId))
        | (_, Group {gChildren}) <- groups
        , Node {nodeId} <- gChildren
        ]
      typeNames =
        [ (stateT, "the state type"), (eventT, "the event type")
        , (transitionT, "the transition type"), (entryT, "the entry type") ]
          ++ [(gType, "compound state " ++ show (T.unpack nodeId)) | (Just Node {nodeId}, Group {gType}) <- groups]
  checkClashes "constructor"
    ( [(transitionT, "the transition type"), (startCon, "the entry type"), (stepCon, "the entry type")]
        ++ stateCons ++ eventCons )
  checkClashes "type" typeNames

  stateDecs <- concat <$> mapM (groupDecs nameFor groupOf . snd) groups
  eventDec <-
    dataD (cxt []) eventT [] Nothing
      [ normalC c [bangType (bang noSourceUnpackedness noSourceStrictness) (conT (mkName (T.unpack f))) | f <- fs]
      | (c, fs, _, _) <- eventInfos
      ]
      [derivClause Nothing (map conT eventDerivs)]
  -- What callbacks receive. A record rather than positional arguments, so a
  -- callback names what it reads and ignores the rest. An exit callback only
  -- runs inside a step, so it gets the record itself; an entry callback may
  -- also run while the chart starts, which leaves nothing and has no event.
  transitionDec <- do
    let field n t = varBangType n (bangType (bang noSourceUnpackedness noSourceStrictness) t)
    dataD (cxt []) transitionT [] Nothing
      [ recC transitionT
          [ field fromField (conT stateT)
          , field eventField (conT eventT)
          , field toField (conT stateT)
          ]
      ]
      [derivClause Nothing (map conT [''Show, ''Read, ''Eq])]
  entryDec <- do
    let plain t = bangType (bang noSourceUnpackedness noSourceStrictness) (conT t)
    dataD (cxt []) entryT [] Nothing
      [normalC startCon [plain stateT], normalC stepCon [plain transitionT]]
      [derivClause Nothing (map conT [''Show, ''Read, ''Eq])]

  let eventNameE
        | null eventInfos = [| \_ -> error "eventName: chart has no events" |]
        | otherwise =
            lamCaseE
              [ match (conP c (replicate (length fs) wildP)) (normalB (lift name)) []
              | (c, fs, _, name) <- eventInfos
              ]
      -- Total for every state the interpreter can complete, which is exactly
      -- the states that got a DoneX constructor.
      doneEventE = do
        sid <- newName "sid"
        lam1E (varP sid) $
          foldr
            (\n rest ->
               [| if $(varE sid) == $(lift (nodeId n))
                    then $(conE (mkName ("Done" ++ T.unpack (nodeId n))))
                    else $rest |])
            [| error ("Statechart: no done.state constructor for " ++ show $(varE sid) ++ "; this is a bug in scxml-statecharts") |]
            completing

  defSig <- sigD defName [t| Def $(conT stateT) $(conT eventT) |]
  defDec <-
    valD (varP defName)
      (normalB
        [| Def
             { defChart = $(lift ch)
             , defEventName = $eventNameE
             , defDoneEvent = $doneEventE
             , defToConfig = $(varE rootTo)
             , defFromConfig = $(varE rootFrom)
             } |])
      []

  -- Hooks dispatching to the callbacks named in the SCXML. Entry callbacks
  -- take an FsmEntry and return m (Maybe FsmEvent), exit callbacks take an
  -- FsmTransition and return m (). Inlined into each generated function
  -- rather than shared, so a polymorphic monad does not hit the monomorphism
  -- restriction.
  -- Underscore-prefixed so that a chart with no callbacks of some phase does
  -- not emit an unused-match warning in the user's module.
  let hooksE = do
        name <- newName "_name"
        cause <- newName "_cause"
        from <- newName "_from"
        ev <- newName "_ev"
        to <- newName "_to"
        let dispatch marker fallback arg fns =
              foldr (\fn rest -> [| if $(varE name) == $(lift fn) then $marker ($(varE (mkName (T.unpack fn))) $arg) else $rest |])
                    fallback fns
            step f e t = [| $(conE transitionT) $f $e $t |]
            entry =
              [| case $(varE cause) of
                   Nothing -> $(conE startCon) $(varE to)
                   Just (f, e) -> $(conE stepCon) $(step [| f |] [| e |] (varE to)) |]
            runEntryE =
              lamE [varP name, varP cause, varP to]
                (dispatch [| Run.entryAction |] [| pure Nothing |] entry entryActions)
            runExitE =
              lamE [varP name, varP from, varP ev, varP to]
                (dispatch [| Run.exitAction |] [| pure () |] (step (varE from) (varE ev) (varE to)) exitActions)
        [| Run.Hooks {Run.runEntry = $runEntryE, Run.runExit = $runExitE} |]
  startDec <- valD (varP startName) (normalB [| Run.start $(varE defName) $hooksE |]) []
  stepDec <- do
    st <- newName "st"
    ev <- newName "ev"
    funD stepName [clause [varP st, varP ev] (normalB [| Run.stepOrStay $(varE defName) $hooksE $(varE st) $(varE ev) |]) []]

  -- Storing a state outside Haskell, as the set of active state ids. Generated
  -- rather than exported so that a chart needs no imports beyond the
  -- quasiquoter itself.
  let toCfg = varE rootTo
      fromCfg = varE rootFrom
  toIdsSig <- sigD toIdsName [t| $(conT stateT) -> [Text] |]
  toIdsDec <- do
    x <- newName "st"
    funD toIdsName [clause [varP x] (normalB [| Set.toAscList ($toCfg $(varE x)) |]) []]
  fromIdsSig <- sigD fromIdsName [t| [Text] -> Maybe $(conT stateT) |]
  fromIdsDec <- do
    ids <- newName "ids"
    given <- newName "given"
    st <- newName "st"
    let body =
          [| let $(varP given) = Set.fromList $(varE ids)
              in case $fromCfg $(varE given) of
                   -- Round-trip so that a set which merely starts like a valid
                   -- one, or is missing part of a configuration, is rejected.
                   Just $(varP st) | $toCfg $(varE st) == $(varE given) -> Just $(varE st)
                   _ -> Nothing |]
    funD fromIdsName [clause [varP ids] (normalB body) []]

  pure (stateDecs ++ [eventDec, transitionDec, entryDec, defSig, defDec, startDec, stepDec, toIdsSig, toIdsDec, fromIdsSig, fromIdsDec])

-- | Data type plus configuration conversions for one compound-like node.
groupDecs :: (StateId -> Name) -> (StateId -> Q Group) -> Group -> Q [Dec]
groupDecs nameFor groupOf Group {gType, gChildren, gTo, gFrom} = do
  shapes <- mapM childShape gChildren
  let dataDec =
        dataD (cxt []) gType [] Nothing
          [normalC con [bangType (bang noSourceUnpackedness noSourceStrictness) (conT f) | f <- fields] | (con, fields, _, _) <- shapes]
          [derivClause Nothing (map conT [''Show, ''Read, ''Eq, ''Ord])]
      toDec = do
        x <- newName "x"
        alts <- forM shapes $ \(con, fields, toE, _) -> do
          vars <- mapM (const (newName "r")) fields
          match (conP con (map varP vars)) (normalB (toE vars)) []
        funD gTo [clause [varP x] (normalB (caseE (varE x) (map pure alts))) []]
      fromDec = do
        cfg <- newName "cfg"
        let body =
              foldr
                (\(Node {nodeId}, (_, _, _, rebuild)) rest -> [| if Set.member $(lift nodeId) $(varE cfg) then $(rebuild cfg) else $rest |])
                [| Nothing |]
                (zip gChildren shapes)
        funD gFrom [clause [varP cfg] (normalB body) []]
  sequence
    [ dataDec
    , sigD gTo [t| $(conT gType) -> Set.Set Text |]
    , toDec
    , sigD gFrom [t| Set.Set Text -> Maybe $(conT gType) |]
    , fromDec
    ]
  where
    -- For one child: (constructor, field types, config-of-fields, rebuild-from-config)
    childShape :: Node -> Q (Name, [Name], [Name] -> Q Exp, Name -> Q Exp)
    -- The groups below are the children's, so their fields are renamed rather
    -- than punned, which would shadow this group's own.
    childShape Node {nodeId = sid, nodeKind} = do
      let con = nameFor sid
      case nodeKind of
        Compound _ -> do
          Group {gType = subType, gTo = subTo, gFrom = subFrom} <- groupOf sid
          pure
            ( con
            , [subType]
            , \vs -> case vs of
                [v] -> [| Set.insert $(lift sid) ($(varE subTo) $(varE v)) |]
                _ -> fail "scxml: internal error, compound state expects exactly one field"
            , \cfg -> [| fmap $(conE con) ($(varE subFrom) $(varE cfg)) |]
            )
        Parallel regionNodes -> do
          regions <- forM (NE.toList regionNodes) $ \Node {nodeId = rid, nodeKind = regionKind} -> case regionKind of
            Compound _ -> Just <$> groupOf rid
            Parallel _ ->
              fail $
                "scxml: <parallel> " ++ T.unpack rid ++ " is directly inside another <parallel>, which is not supported."
                  ++ " Flatten it: its regions can become regions of the outer <parallel>, since all of them are active at once either way."
                  ++ " The only thing flattening loses is a done.state event for the inner one on its own."
            _ -> pure Nothing
          let regionIds = map nodeId (NE.toList regionNodes)
              fieldTypes = [subType | Just Group {gType = subType} <- regions]
              toE vs =
                let go [] _ = []
                    go ((r, Nothing) : rest) vars = [| Set.singleton $(lift r) |] : go rest vars
                    go ((r, Just Group {gTo = subTo}) : rest) (v : vars) = [| Set.insert $(lift r) ($(varE subTo) $(varE v)) |] : go rest vars
                    go _ [] = error "scxml: internal error, region/field mismatch"
                 in [| Set.insert $(lift sid) (Set.unions $(listE (go (zip regionIds regions) vs))) |]
              rebuild cfg = foldl (\acc subFrom -> [| $acc <*> $(varE subFrom) $(varE cfg) |]) [| pure $(conE con) |] [subFrom | Just Group {gFrom = subFrom} <- regions]
          pure (con, fieldTypes, toE, rebuild)
        _ ->
          pure
            ( con
            , []
            , \_ -> [| Set.singleton $(lift sid) |]
            , \_ -> [| Just $(conE con) |]
            )

checkClashes :: String -> [(Name, String)] -> Q ()
checkClashes what named = do
  let byName = Map.fromListWith (++) [(nameBase n, [origin]) | (n, origin) <- named]
      clashes = [(n, os) | (n, os) <- Map.toList byName, length os > 1]
  unless (null clashes) $
    fail $ unlines $
      ("scxml: generated " ++ what ++ " names clash; rename one of the SCXML identifiers:")
        : ["  " ++ n ++ " would be generated for " ++ commaList (reverse os) | (n, os) <- clashes]
  where
    commaList [a, b] = a ++ " and " ++ b
    commaList xs = foldr1 (\a b -> a ++ ", " ++ b) xs

-- | Callback names in @<script>@ must be Haskell variable names, optionally
-- qualified (@Inventory.reserve@).
checkVarName :: Text -> Either String ()
checkVarName raw
  | ok = Right ()
  | otherwise = Left ("<script> callback " ++ show (T.unpack raw) ++ " is not a Haskell function name (expected something like reserveStock or Inventory.reserve)")
  where
    segments = T.splitOn (T.pack ".") raw
    ok = not (null segments) && all isModulePart (init segments) && isVar (last segments)
    isVar t = case T.unpack t of
      c : cs -> (isLower c || c == '_') && all (\x -> isAlphaNum x || x == '_' || x == '\'') cs
      [] -> False
    isModulePart t = case T.unpack t of
      c : cs -> isUpper c && all (\x -> isAlphaNum x || x == '_' || x == '\'') cs
      [] -> False

orFail :: Either String a -> Q a
orFail = either (fail . ("scxml: " ++)) pure