Free Python-style generators

FreeGenerators.hs

module FreeGenerators where

import Control.Monad.Trans
import Control.Monad.Trans.Free

data YieldF i o r = YieldF o (i -> r)

instance Functor (YieldF i o) where
  fmap f (YieldF o k) = YieldF o (f . k)

type Generator i o = FreeT (YieldF i o)
data X
type Effect = Generator X X

yield :: Monad m => o -> Generator i o m i
yield o = liftF (YieldF o id)

runEffect :: Monad m => Effect m r -> m r
runEffect (FreeT m) = do
  f <- m
  case f of
    Pure r -> return r
    Free (YieldF o k) -> undefined -- impossible

for :: Monad m => Generator i o m r -> (o -> Generator i' o' m i) -> Generator i' o' m r
for (FreeT m) l = FreeT $ do
  f <- m
  case f of
    Pure r -> return (Pure r)
    Free (YieldF o k) -> runFreeT $ l o >>= (`for` l) . k

gen :: Monad m => Generator Int String m Int
gen = do
  i <- yield "i"
  j <- yield "j"
  return (i + j)

example :: IO Int
example = runEffect $ for gen $ \s -> do
  lift $ putStrLn ("Got " ++ s)
  return 123