FreerMonadFx.hs

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

import Data.Bifunctor (first)
import Data.Functor.Sum
import Data.Functor.Coyoneda
import Control.Monad
import Control.Monad.Free
import Control.Monad.Identity

import qualified Control.Monad.Except as E
import qualified Control.Monad.State as S
import qualified Data.Map.Strict as M
import qualified Data.HashSet as S
import qualified Data.Tree as T

data Union (r :: [*])

type Freer f g = Free (Coyoneda (Sum f g))

runFX :: Monad m => (forall x. f x -> m x) -> (forall x. g x -> m x) -> Freer f g a -> m a
runFX _ _ (Pure a) = pure a
runFX interF interG (Free (Coyoneda f (InL fb))) =
  let mb = interF fb
  in mb >>= \b -> runFX interF interG $ f b
runFX interF interG (Free (Coyoneda f (InR gb))) =
  let mb = interG gb
  in mb >>= \b -> runFX interF interG $ f b

etaF :: f a -> Freer f g a
etaF fa = Free $ Coyoneda Pure (InL fa)

etaG :: g a -> Freer f g a
etaG ga = Free $ Coyoneda Pure (InR ga)

data State s a where
  Get :: State s s
  Put :: s -> State s ()

unState :: State s a -> (s -> (a, s))
unState Get s = (s, s)
unState (Put s) _ = ((), s)

getF :: Freer (State s) g s
getF = etaF Get

putF :: s -> Freer (State s) g ()
putF = etaF . Put

data Except e a where
  ThrowError :: e -> Except e ()

throwErrorF :: e -> Freer f (Except e) ()
throwErrorF = etaG . 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 -> Freer (State (M.Map VariableName Int)) (Except String) (AST Int)
assignIndexToVariables ast variables = forM ast $ \var -> do
  unless (var `S.member` variables) $
    etaG $ ThrowError $ "Unknown Variable " <> var
  cache <- etaF Get
  case M.lookup var cache of
    Just index -> pure index
    Nothing -> do
      let index = M.size cache
      etaF $ Put $ M.insert var index cache
      pure index

runApp :: Freer (State s) (Except e) a -> E.ExceptT e (S.State s) a
runApp = runFX (E.ExceptT . S.state . (\f -> first Right . f) . unState) (\(ThrowError e) -> E.throwError e)

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 S.evalState mempty $ E.runExceptT $ runApp $ assignIndexToVariables ast vars