Iterator.hs

{-# LANGUAGE FlexibleInstances, FunctionalDependencies, MultiParamTypeClasses #-}

import Control.Monad
import Control.Monad.Loops
import Control.Monad.State
import Data.Char

class Iterator a b r | a -> b r where
    next :: (Monad m) => b -> StateT a m (Maybe r)

data RangeIterator a = RangeIterator { rangeCurrent :: a, rangeEnd :: a }
    deriving (Show)

rangeIterator :: (Eq a, Enum a, Ord a) => a -> a -> RangeIterator a
rangeIterator = RangeIterator

instance (Eq a, Enum a, Ord a) => Iterator (RangeIterator a) () a where
    next _ = do
        RangeIterator n e <- get
        if n > e then
            return Nothing
        else do
            modify $ \s -> s { rangeCurrent = succ n }
            return $ Just n

newtype ListIterator a = ListIterator [a]
    deriving (Show)

instance Iterator (ListIterator a) () a where
    next _ = do
        ListIterator xs <- get
        case xs of
            [] -> return Nothing
            (y:ys) -> put (ListIterator ys) >> return (Just y)

data YellingIterator = YellingIterator 
    deriving (Show)

instance Iterator YellingIterator String String where
    next "" = return Nothing
    next s = return . Just $ map toUpper s ++ "!"

runIterator it = flip execStateT it

loopIterator it f = flip evalStateT it $ whileJust (next ()) f

main = do
    putStrLn "Enter starting number:"
    -- Notice that readLn is inferring the type at compile time!
    start <- readLn
    putStrLn "Enter ending number:"
    end <- readLn
    let it = rangeIterator start end

    putStrLn "\nLet's manually use next a couple times..."
    runIterator it $ do
        n <- next ()
        liftIO . putStrLn $ "Got one: " ++ show n
        n <- next ()
        liftIO . putStrLn $ "Here's another: " ++ show n

    putStrLn "\nNow let's crappily implement fizzbuzz..."
    loopIterator it $ \n -> do
        let fizzbuzz = (if n `mod` 3 == 0 then "Fizz" else "") ++
                       (if n `mod` 5 == 0 then "Buzz" else "")
        let output = if fizzbuzz == "" then show n else fizzbuzz
        liftIO $ putStr (show n ++ ":\t") >> putStrLn output

    putStr "\nNumber of elements in our iterator: "
    print =<< liftM length (loopIterator it return)

    let it' = rangeIterator 'a' 'z'
    putStrLn "\nWe can build a RangeIterator from anything supporting Eq, Ord, and Enum"
    putStrLn $ "For instance, let's loop through " ++ show it'
    print =<< loopIterator it' return

    let listIt = ListIterator [1..10]
    putStrLn $ "\nWe can also implement other iterators, such as " ++ show listIt
    loopIterator listIt (liftIO . print)

    putStrLn "\nLet's have fun with an iterator that handles next arguments"
    putStrLn "Input text to the YellingIterator:"
    runIterator YellingIterator $ (whileJust_ (next =<< liftIO getLine)) (liftIO . putStrLn)
    return ()