Iteratees part II

I'd like to go over the basics of the enumerator package for Haskell, covering the types, and re-solving the cube-permutation problem from the last post. I'm not the first to cover this, see this explanation of the enumerator package.

The first datatype defined in the enumerator package is Stream:

data Stream a

= Chunks [a]

| EOF

deriving (Show, Eq)

Which is fairly straightforward. Next is a Step:



data Step a m b
= Continue (Stream a -> Iteratee a m b)
| Yield b (Stream a)
| Error Exc.SomeException



I'll explain what an Iteratee is next. A step is one of 3 states: Continue, it can accept more data. Yield, it can't accept any more data, and Error, an error occured. Note that unlike EarlyFoldable, a Step doesn't have a type for the state, instead it passes a continuation that holds the state.


So what is an Iteratee? It's simply a Step wrapped in a Monad.

newtype Iteratee a m b = Iteratee

{ runIteratee :: m (Step a m b)

}

Let's see if we can solve the cube problem using an iteratee.




module Main where


import Data.List
import Data.Enumerator hiding (head, foldl', repeat)
import qualified Data.Enumerator.List as EL (head)
import qualified Data.Map as M


digitize :: (Integral a) => a -> [a]
digitize 0 = [0]
digitize n = (reverse . (unfoldr strip_digit)) n
    where 
        strip_digit 0 = Nothing
        strip_digit a =
            case (a `divMod` 10) of
                (a,b)  -> Just (b,a)


undigitize :: Integral a => [a] -> a
undigitize = foldl' (\x y -> x*10+y) 0


--digitize and undigitize are for number processing


cubes :: [Integer]
cubes = scanl (+) 1 $ scanl (+) 7 $ scanl (+) 12 $ repeat 6


make_key :: Integer -> Integer
make_key = undigitize . sortBy (flip compare) . digitize


update_map :: Integer -> M.Map Integer (Integer, Int) -> M.Map Integer (Integer, Int)
update_map n = M.insertWith
  (\(n2, count2) (n1, count1) -> (n1, (count1+count2)))
  (make_key n) (n, 1)


--cubes, make_key and update_map are the same as from last time.


cube_perm_iteratee :: (Monad m) => Iteratee Integer m Integer
cube_perm_iteratee = go M.empty
  where
    go m = do -- Iteratee's are a Monad, to be covered later.
      (Just n) <- EL.head --EL.head reads at most one element.
      case update_map n m of
        m' | (snd $ m' M.! (make_key n)) == 5 ->
          yield (fst $ m' M.! (make_key n)) (Chunks [])
        m' -> go m'


main = do
  -- The driver code will be covered on another article on Iteratees
  result <- run_ $ (enumList 100 cubes) $$ cube_perm_iteratee
  print result