Generators in Haskell.

Generators.hs

{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

module Generators where

import Prelude hiding (take, filter, foldr)

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

import Data.Foldable
import Data.Monoid
import Data.Traversable
import Data.Void

data Generator i o m r
  = Done r
  | Step (m (Generator i o m r))
  | Yield o (i -> Generator i o m r)

instance Monad m => Functor (Generator i o m) where
  fmap f (Done r) = Done (f r)
  fmap f (Step m) = Step (liftM (fmap f) m)
  fmap f (Yield o k) = Yield o (fmap f . k)

instance Monad m => Applicative (Generator i o m) where
  pure = Done
  Done f <*> g = fmap f g
  Step m <*> g = Step (liftM (<*> g) m)
  Yield o k <*> g = Yield o ((<*> g) . k)

instance Monad m => Monad (Generator i o m) where
  return = Done
  Done a >>= f = f a
  Step m >>= f = Step (liftM (>>= f) m)
  Yield o k >>= f = Yield o ((>>= f) . k)

instance (Monad m, Monoid r) => Monoid (Generator i o m r) where
  mempty = return mempty
  mappend = liftA2 mappend

instance MonadTrans (Generator i o) where
  lift = Step . liftM Done

pipe :: Monad m => Generator i x m r -> Generator x o m r -> Generator i o m r
pipe g (Done r) = Done r
pipe g (Step mg') = Step (liftM (pipe g) mg')
pipe (Done r) (Yield o k') = Yield o (\i -> Done r)
pipe (Step m) (Yield o k') = Step (liftM (`pipe` (Yield o k')) m)
pipe (Yield x k) (Yield o k') = Yield o (\i -> pipe (k i) (k' x))

g >-> g' = pipe g g'
infixl 7 >->

yield :: Monad m => o -> Generator i o m i
yield o = Yield o Done

type Effect m r = Generator Void Void m r

runEffect :: Monad m => Effect m r -> m r
runEffect (Done r) = return r
runEffect (Step mg) = mg >>= runEffect
runEffect (Yield o k) = absurd o -- impossible

next :: Monad m => Generator i o m r -> m (Either r (o, i -> Generator i o m r))
next (Done r) = return (Left r)
next (Step m) = m >>= next
next (Yield o k) = return (Right (o, k))

each :: Monad m => Generator i o m r -> (o -> Generator i' o' m i) -> Generator i' o' m r
each (Done r) l = Done r
each (Step mg) l = Step (liftM (`each` l) mg)
each (Yield o k) l = l o >>= (`each` l) . k

filter :: Monad m => (o -> Bool) -> Generator () o m r -> Generator () o m r
filter p g = each g $ \o -> if (p o) then yield o else return ()

from :: (Monad m, Foldable f) => f a -> Generator i a m ()
from = foldr ((>>) . yield) (return ())

take :: Monad m => Int -> Generator i o m r -> Generator i o m ()
take n g | n <= 0 = Done ()
take n (Done r) = Done ()
take n (Step mg) = Step (liftM (take n) mg)
take n (Yield o k) = Yield o (take (n - 1) . k)

fibs :: Monad m => Generator () Int m ()
fibs = go 0 1
  where
    go i j = yield i >> go j (i + j)

stdinLn = do
  l <- lift getLine
  yield l
  stdinLn

example = runEffect $ each (take 10 fibs) $ \i ->
  lift $ putStrLn (show i)

ints :: Generator () Int IO ()
ints = go 0
  where
    go i = do
      yield i
      go (i + 1)

inMap :: Monad m => Generator i o m r -> (i' -> i) -> Generator i' o m r
inMap g f = each g $ \o -> do
  i <- yield o
  return (f i)

outMap :: Monad m => Generator i o m r -> (o -> o') -> Generator i o' m r
outMap g f = each g $ \o -> do
  yield (f o)

printInts = each ints $ \i ->
  lift (putStrLn ("We got " ++ show i))

doublingFibs = each fibs $ \i -> do
  yield i
  yield i

example2 = each (take 5 stdinLn) (lift . putStrLn)

example3 = each (take 2 fibs) $ \n -> do
  lift $ putStrLn $ "Current fib = " ++ show n
  each (take 2 stdinLn) $ \s ->
    yield (n, s)

chattyGen = do
  lift (putStrLn "hi!")
  lift (putStrLn "hi!")
  lift (putStrLn "hi!")
  lift (putStrLn "hi!")
  yield "foo"

anotherGenerator i = do
  yield (i + 1)
  yield (i + 2)
  yield (i + 3)

nestingExample i = do
  yield i
  anotherGenerator i
  yield (i + 10)

myGen x = do
  i <- yield x
  j <- yield (x + 1)
  k <- yield (x + 2)
  return (i, j, k)

myOtherGen = each (myGen 100) $ \i -> do
  lift $ putStrLn ("Inside myOtherGen!! " ++ show i)
  return (show i)

main = do
  Right (i, k) <- next fibs
  putStrLn (show i)
  Right (i', k') <- next (k ())
  putStrLn (show i')
  Right (i'', k'') <- next (k' ())
  putStrLn (show i'')