Sunday, April 17, 2011

Implementing scanl for Iteratees

In looking over Data.Enumerator.List, I noticed that scanl was missing, So I set out to write it. I'm posting it as a literate Haskell file. I intend to create a patch for the maintainer.

Start with a standard module declaration


>module Data.Enumerator.Scanl where
>
>import Prelude hiding (scanl, scanl1)
>import Data.Enumerator
>import qualified Data.Enumerator.List as EL


The code seems strangely devoid of an operator for repeatedly chaining enumeratees together, so I created one.


>infixl 0 $|
>($|) :: Monad m
>      => Enumerator ao m (Step ai m b)
>      -> Enumeratee ao ai m b
>      -> Enumerator ai m b
>($|) = joinE


scanl takes a function and a seed value, and combines the seed and any input to get a new seed. The seeds (including the first) are fed along the pipe.


>scanl :: (Monad m) => (ai -> ao -> ai) -> ai -> Enumeratee ao ai m b


I borrowed the structure of Data.Enumerator.List.map for most of this. CheckDone is a common construction for Enumeratees, if the inner iteratee yields, the outer one yields as well.


>scanl op ai0 = checkDone (continue . step ai0) where


The final step here is important. For map it is:


  step k EOF = yield (Continue k) EOF


For scanl, we have one more piece of data (the last seed) that needs to be fed into the iteratee.


>  step ai0 k EOF = k (Chunks [ai0]) >>== \step -> yield step EOF
>  step ai0 k (Chunks xs) = loop ai0 k xs
>
>  loop ai0 k [] = continue (step ai0 k)


Note that we compute the next seed and store it, but send the old seed to the iteratee.


>  loop ai0 k (x:xs) = case ai0 `op` x of
>    ai1 -> k (Chunks [ai0]) >>==
>      checkDoneEx (Chunks xs) (\k' -> loop ai1 k' xs)


And now, some test code:


>cubes :: Monad m => Enumerator Integer m b
>cubes = EL.replicate 7 6 $| scanl (+) 12 $| scanl (+) 7 $| scanl (+) 1


The important thing about the test code is that there should be exactly 10 cubes in the list. Each scanl adds exactly one element, and we start with 7 repetitions of 6, so we should (and do) have 10 cubes.
and to test:

>main = do
>  list <- run_ $ cubes $$ EL.take 10
>  print list

No comments:

Post a Comment