Free Monad Effects example

FreeMonadFx.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
module FreeMonadFx where

import Data.Functor.Sum
import Control.Monad
import Control.Monad.Except
import Control.Monad.Free
import Control.Monad.State

import qualified Data.Map.Strict as M
import qualified Data.HashSet as S
import qualified Data.Tree as T

type DerpFX f g = Free (Sum f g)

runDerpFX :: s -> DerpFX (State s) (Either e) a -> Either e (a, s)
runDerpFX s (Pure a) = Right (a, s)
runDerpFX s (Free (InL m)) = let (m', s') = runState m s in runDerpFX s' m'
runDerpFX s (Free (InR (Left e))) = throwError e
runDerpFX s (Free (InR (Right m))) = runDerpFX s m

runFX' :: Monad m => (forall x. f x -> m x) -> (forall x. g x -> m x) -> DerpFX f g a -> m a
runFX' _ _ (Pure a) = pure a
runFX' interF interG (Free (InL f)) = let m = interF f in m >>= runFX' interF interG
runFX' interF interG (Free (InR g)) = let m = interG g in m >>= runFX' interF interG

runApp :: Monoid s => DerpFX (State s) (Either e) a -> ExceptT e (State s) a
runApp = runFX' lift (ExceptT . pure)

interState :: Monad m => s -> Free (State s) a -> m a
interState s (Pure a) = pure a
interState s (Free m) = let (free', s') = runState m s in interState s' free'

interEither :: Monad m => Free (Either e) a -> m (Either e a)
interEither (Pure a) = pure (Right a)
interEither (Free (Left e)) = pure (Left e)
interEither (Free (Right a)) = interEither a

eta :: Functor f => f a -> Free f a
eta = Free . fmap Pure

type FState s = Free (State s)

getF :: FState s s
getF = eta get

putF :: s -> FState s ()
putF = eta . put

runStateF :: FState s a -> s -> (a, s)
runStateF (Pure x) s = (x, s)
runStateF (Free m) s =
  let (m', s') = runState m s in runStateF m' s'

type FExcept e m = Free (ExceptT e m)

throwErrorF :: Monad m => e -> FExcept e m a
throwErrorF = eta . throwError

type VariableName = String
type Variables = S.HashSet VariableName

data AST a = Leaf a | Node (AST a) (AST a)
  deriving (Show, Functor, Foldable, Traversable)

assignIndexToVariables :: AST VariableName -> Variables -> DerpFX (State (M.Map VariableName Int)) (Either String) (AST Int)
assignIndexToVariables ast variables = forM ast $ \var -> do
  unless (var `S.member` variables) $
    eta $ InR $ throwError $ "Unknown Variable " <> var
  cache <- eta $ InL get
  case M.lookup var cache of
    Just index -> pure index
    Nothing -> do
      let index = M.size cache
      eta $ InL $ put $ M.insert var index cache
      pure index

main :: IO ()
main =
  let vars = S.fromList ["a", "b", "c"]
      ast = Node (Leaf "a") (Node (Leaf "b") (Node (Leaf "a") (Leaf "c")))
  in print $ flip evalState mempty $ runExceptT $ runApp $ assignIndexToVariables ast vars