[Haskell-cafe] sequence and sequence_ for iteratees

[Maciej Wos]

Hello Cafe!

I have written some iteratee functions that I found to be very very
useful and I hope they will soon make it to the iteratee library.
However, I'd like to get some feedback first, particularly about the
error handling.

I'm sending this here rather than to iteratee mailing list at
projects.haskell.org because the latter is (and has been for a while!)
down...

Anyway, enumSequence and enumSequence_ are inspired by Prelude's
sequence and sequence_. They are useful when one has to deal with
several iteratees consuming the same data. For instance, one could use
enumSequence as follows:

> run $ joinIM $ enumPureNChunk [1..100] 3 $
    enumSequence [I.head, I.head >>= \x -> I.head >>= \y -> return
(x+y), I.last]

which produces:

[1,3,100]

Each iteratee in the list is given the same input stream. Also,
enumSequence consumes as much of the stream as the iteratee in the
list that consumes the most. In the above example there is no stream
left after enumSequence finishes because I.last consumes everything.

As an another example, enumSequence below consumes only the first two
elements of the stream and the remainder is passed to stream2list:

> run $ joinIM $ enumPureNChunk [1..10] 3 $
    (enumSequence [I.head, I.head >> I.head] >> stream2list)

[3,4,5,6,7,8,9,10]

The code for enumSequence is enclosed below (enumSequence_ is almost
identical!) To make it complete though I should add some sort of error
handling. I'm not quite sure however what would be the best thing to
do. For instance, what should happen if the stream is finished, but
one of the iteratees is not done yet? Should the whole enumSequence
fail? Similarly, should the whole enumSequence fail if one of the
iteratees throws an error?

I guess throwing some sort of recoverable error could work. But I
still need to figure out how to do that!

-- Maciej

########## code ##########

enumSequence :: forall m s a el . (Monad m, LL.ListLike s el, Nullable s)
             => [Iteratee s m a]
             -> Iteratee s m [a]
enumSequence is = liftI step
    where
        step :: Stream s -> Iteratee s m [a]
        step s@(Chunk xs) | LL.null xs = liftI step
                          | otherwise  = do
              let is'  = map (joinIM . enumPure1Chunk xs) is
              allDone <- lift (checkIfDone is')
              if allDone then uncurry idone =<< lift (collectResults is')
                         else enumSequence (updateChunk s is')
        -- TODO: should return an error if not all iteratees are done
        step (EOF _) = uncurry idone =<< lift (collectResults . map
(joinIM . enumEof) $ is)

        -- returns true if *all* iteratees are done; otherwise returns false
        checkIfDone :: [Iteratee s m a] -> m Bool
        checkIfDone = liftM and . mapM (\i -> runIter i onDone onCont)
            where
                onDone _ _ = return True
                onCont _ _ = return False

        -- returns a list of result values and the unconsumed part of the stream
        collectResults :: [Iteratee s m a] -> m ([a], Stream s)
        collectResults = liftM (id *** foldl1 shortest)
                         . mapAndUnzipM (\i -> runIter i onDone onCont)
            where
                onDone a s = return (a,s)
                onCont _ _ = error "enumSequence: collectResults
failed; all iteratees should be done"

        shortest :: Stream s -> Stream s -> Stream s
        shortest (Chunk xs) (Chunk ys)
            | LL.length xs > LL.length ys = Chunk ys
            | otherwise                   = Chunk xs
        shortest s@(EOF _) _ = s
        shortest _ s@(EOF _) = s


        -- iteratee in *done* state holds the unconsumed part of the chunk it
        -- was given; this chunk needs to be discarded when we move further in
        -- the stream
        updateChunk :: Stream s -> [Iteratee s m a] -> [Iteratee s m a]
        updateChunk s = map (\i -> joinIM $ runIter i (\a _ -> return
$ idone a s) icontM)