ListT done right

ListT.hs

{-# LANGUAGE TypeFamilies    #-}
{-# LANGUAGE DeriveFunctor   #-}
{-# LANGUAGE OverloadedLists #-}

module ListT where

import Control.Applicative
import Control.Monad
import Control.Monad.Trans
import Data.Monoid
import GHC.Exts

newtype ListT m a =
  ListT { next :: m (Maybe (a, ListT m a)) }

--------------------------------------------------------------------------------

nil :: Monad m => m a -> ListT m b
nil act = ListT (liftM (const Nothing) act)

cons :: Monad m => m a -> ListT m a -> ListT m a
cons act tail = ListT (liftM (\a -> Just (a, tail)) act)

gen :: Monad m => m a -> ListT m a
gen m = z where z = ListT (liftM (\a -> Just (a, z)) m)

-- This will blow stacks
stream :: Monad m => ListT m a -> m [a]
stream l = do
  m <- next l
  case m of
    Nothing      -> return []
    Just (a, l') -> do
      as <- stream l'
      return (a:as)

--------------------------------------------------------------------------------

instance Monad m => Functor (ListT m) where
  fmap f l = ListT $ do
    m <- next l
    return $ case m of
      Nothing      -> Nothing
      Just (a, l') -> Just (f a, fmap f l')

instance Monad m => Monoid (ListT m a) where
  mempty      = nil (return ())
  mappend a b = ListT $ do
    m <- next a
    case m of
      Nothing      -> next b
      Just (x, a') -> return $ Just (x, mappend a' b)

-- Kind of shoddy
instance Monad m => IsList (ListT m a) where
  type Item (ListT m a) = a
  fromList l = case l of
    []     -> mempty
    a : as -> cons (return a) (fromList as)
  toList = error "nope!"

instance Monad m => Applicative (ListT m) where
  pure  = return
  (<*>) = ap

instance Monad m => Monad (ListT m) where
  return a = cons (return a) (nil (return ()))
  l >>= f = ListT $ do
    m <- next l
    case m of
      Nothing      -> return Nothing
      Just (a, l') -> next (mappend (f a) (l' >>= f))

instance MonadTrans ListT where
  lift m = ListT $ m >>= next . return

--------------------------------------------------------------------------------

-- A solution to Joachim Breitner's "Constructing a list in a Monad"
-- <http://www.joachim-breitner.de/blog/620-Constructing_a_list_in_a_Monad>

getInput :: IO Int
getInput = liftM read getLine

getn :: Int -> ListT IO Int
getn 0 = mempty
getn n = lift getInput <> getn (n - 1)