A monad transformer for generator coroutines

Generator.hs

-- | A monad transformer for generator coroutines
module Control.Monad.Generator where

import Control.Applicative
import Control.Monad
import Control.Monad.Trans

-- | The type of generators yielding intermediate results of type @t@ in base
-- monad @m@, with a final result of type @r@
data Generator t m r = Generator {runGen :: m (GeneratorState t m r)}

instance Functor m => Functor (Generator t m) where
  fmap f (Generator m) = Generator $ fmap (fmap f) m

instance (Functor m, Monad m) => Applicative (Generator t m) where pure = return
                                                                   (<*>) = ap

instance Monad m => Monad (Generator t m) where
  return = Generator . return . Done
  Generator m >>= f = Generator $ do
    state <- m
    case state of Done r -> runGen $ f r
                  More t g -> return $ More t $ g >>= f

instance MonadTrans (Generator t) where lift m = Generator $ m >>= return . Done

instance (Functor m, MonadIO m) => MonadIO (Generator t m) where
  liftIO = Generator . fmap Done . liftIO

data GeneratorState t m r
  -- | The suspended state, with an intermediate result value and a
  -- continuation generator
  = More t (Generator t m r)
  -- | The completed state with a final result value
  | Done r

instance Functor m => Functor (GeneratorState t m) where
  fmap f (More t gen) = More t $ fmap f gen
  fmap f (Done r) = Done $ f r

-- | Suspend the generator, producing the argument as an intermediate result
yield :: (Functor m, Monad m) => t -> Generator t m ()
yield x = Generator $ return $ More x $ return ()

-- | Run a generator to completion, printing the results to stdout as they are
-- produced
printValues :: Show t => Generator t IO r -> IO r
printValues = genState id $ \t g -> print t >> printValues g

-- | Run a generator to completion, producing a list of the results
listValues :: (Functor m, Monad m) => Generator t m r -> m [t]
listValues = genState (const []) (\t g -> fmap (t:) $ listValues g)

-- | Run a generator up to the first yield, producing @Just@ the first value
--
-- If the generator never yields, produce @Nothing@.
firstValue :: Monad m => Generator t m r -> m (Maybe t)
firstValue = genState (const Nothing) (\t _ -> return $ Just t)

-- | Construct a generator from a list of values to be yielded
fromList :: Monad m => [t] -> Generator t m ()
fromList [] = Generator $ return $ Done ()
fromList (x:xs) = Generator $ return $ More x $ fromList xs

-- | A generic case analysis on generator states
genState :: Monad m
         => (r -> a) -- ^ Handler for the final result
         -> (t -> Generator t m r -> m a) -- ^ Handler for the suspended case
         -> Generator t m r
         -> m a
genState done more gen = do state <- runGen gen
                            case state of Done r -> return $ done r
                                          More t g' -> more t g'