Algebraic effects implemented using threads

AlgEffWithThreads.hs

-- Algebraic effects using threads

{-# LANGUAGE BlockArguments, GADTs, RankNTypes, EmptyCase, LambdaCase #-}
import Control.Concurrent
import Control.Applicative
import Control.Monad

-- Free monad as a threaded computation
-- The `forall` guarantees that `Return` can't be emitted in the middle of
-- a computation, only at the end, in a wrapper defined by `interpret`.
newtype Free f a = Free (forall x. MVar (Event f x) -> IO a)

-- Internal type: a thread either "returns" a final value,
-- or "performs" an operation and waits for its result in the provided `MVar`.
data Event f a where
  Return :: a -> Event f a
  Perform :: f b -> MVar b -> Event f a

-- Internal unwrapper
unFree :: MVar (Event f x) -> Free f a -> IO a
unFree fvar (Free m) = m fvar

-- * Public interface

-- ** Monad

pureFree :: a -> Free f a
pureFree x = Free \_ -> pure x

bindFree :: Free f a -> (a -> Free f b) -> Free f b
bindFree m k = Free \fvar -> unFree fvar m >>= \x -> unFree fvar (k x)

instance Monad (Free f) where
  (>>=) = bindFree

instance Applicative (Free f) where
  pure = pureFree
  (<*>) = ap

instance Functor (Free f) where
  fmap = liftA

-- ** send, runFree and interpret

-- Send an operation `f a` and wait for its result `a`.
send :: f a -> Free f a
send op = Free \fvar -> do
  fvar' <- newEmptyMVar
  putMVar fvar (Perform op fvar')
  takeMVar fvar'

data Empty a

-- Run a pure computation.
-- This is supposed to be `Free Empty a -> a` but I don't want to use `unsafePerformIO` for this self-contained example.
runFree :: Free Empty a -> IO a
runFree (Free m) = do
  evar <- newEmptyMVar
  m evar  -- the mvar should never be filled

-- Interpret a computation with a handler.
-- Continuations (`x -> Free g a`) are one-shot: calling them more than once is an error.
interpret :: (forall x. f x -> (x -> Free g a) -> Free g a) -> Free f a -> Free g a
interpret handler (Free m) = Free \gvar0 -> do
  fvar <- newEmptyMVar
  -- Fork the handled computation
  forkIO do
    x <- m fvar
    putMVar fvar (Return x)
  -- Handle the computation by processing its events.
  let loop = Free \gvar -> takeMVar fvar >>= \case
        Return x -> pure x
        Perform op opvar -> do
          -- The handler may call its continuation at most once.
          callOnce <- newEmptyMVar
          unFree gvar (handler op \y -> Free \gvar' -> do
            tryPutMVar callOnce () >>= \isFirstCall -> when (not isFirstCall) (error "Continuation may be called at most once")
            putMVar opvar y
            unFree gvar' loop)
  unFree gvar0 loop

-- * Example: coroutines

-- Free monad transformer as an ADT (aka. initial encoding)
-- https://hackage.haskell.org/package/free-5.2/docs/Control-Monad-Trans-Free.html#t:FreeT
type ADTFreeT f m a = m (ADTFreeTF f m a)

data ADTFreeTF f m a where
  Pure :: a -> ADTFreeTF f m a
  Send :: f b -> (b -> m (ADTFreeTF f m a)) -> ADTFreeTF f m a

-- From Free to ADTFreeT
toADTFreeT :: Free f a -> ADTFreeT f (Free g) a
toADTFreeT m = interpret handler (fmap Pure m)
  where
    handler :: f x -> (x -> ADTFreeT f (Free g) a) -> ADTFreeT f (Free g) a
    handler op k = pure (Send op k)

-- Coroutine operation: call with argument `a` and result `b`
data Call a b r where
  Call :: a -> Call a b b

-- Interleave two coroutines that call each other.
-- Result has no effect.
coroutine :: Free (Call a b) r -> (a -> Free (Call b a) r) -> Free f r
coroutine m n = toADTFreeT m >>= interleave (toADTFreeT . n)
  where
    interleave :: (a -> ADTFreeT (Call b a) (Free f) r) -> ADTFreeTF (Call a b) (Free f) r -> Free f r
    interleave n (Pure r) = pure r
    interleave n (Send (Call x) k) = n x >>= interleave k

-- Pong coroutine: append "pong" to every input message.
pong :: String -> Free (Call String String) a
pong msg = send (Call (msg ++ " pong")) >>= pong

-- Ping coroutine: make 3 calls, appending "ping" to every intermediate message.
ping :: Free (Call String String) String
ping = do
  x <- send (Call "ping")
  y <- send (Call (x ++ " ping"))
  z <- send (Call (y ++ " ping"))
  pure z

-- Interleave ping and pong.
runExample :: Free Empty String
runExample = coroutine ping pong
-- "pingpongpingpongpingpong"

main :: IO ()
main = runFree runExample >>= print