arrow-like functions for enumeratees

EteeArr.hs

{-# LANGUAGE RankNTypes, ScopedTypeVariables, NoMonomorphismRestriction,
   TupleSections #-}

-- | A couple arrow-like functions for enumeratees.  Consider this example:
--
-- Suppose you have the following
--
-- 1.  Consumer :: Iteratee [(Int,Char)] m a
-- 2.  Mapper   :: Enumeratee Int Char m a
-- 3.  Source   :: Enumerator Int m a
--
-- You would like to use the mapper function to convert the data stream from
-- 'Source' into a tuple of (Int,Char).  With arrows, this would be something
-- like
-- arr (\a -> (a,a)) >>> second mapper
--
-- This module enables something similar for enumeratees
--
-- > *Main Data.Iteratee> let consumer = stream2list :: Monad m=> I [(Int,Char)] m [(Int,Char)]
-- > *Main Data.Iteratee> let mapper = mapChunks (map chr)
-- > *Main Data.Iteratee> let source = enumPureNChunk [1..20] 3 :: Monad m => Enumerator [Int] m a
-- > *Main Data.Iteratee> print =<< run =<< source (joinI $ dupE ><> secondE mapper $ consumer)
-- [(1,'\SOH'),(2,'\STX'),(3,'\ETX'),(4,'\EOT'),(5,'\ENQ'),(6,'\ACK'),(7,'\a'),(8,'\b'),(9,'\t'),(10,'\n'),(11,'\v'),(12,'\f'),(13,'\r'),(14,'\SO'),(15,'\SI'),(16,'\DLE'),(17,'\DC1'),(18,'\DC2'),(19,'\DC3'),(20,'\DC4')]
-- > *Main Data.Iteratee> 
--

module EteeArr
 ( secondE
 , dupE
 ) where

import Data.Iteratee as I
import Control.Monad.Trans
import Control.Monad
import Control.Exception

import Debug.Trace
import Data.Maybe
import Data.ListLike (ListLike)
import Data.Char


-- this is an overconstrained type, maybe the original would be more useful?
secondE
  :: (Monad m)
  => (I [b] (I [a] m) x -> I [b2] (I [a2] m) (I [b] (I [a] m) x))
  -> I [(a, b)] m x
  -> Iteratee [(a2, b2)] m (Iteratee [(a, b)] m x)
secondE mapper = liftM combine . combine . mapper . split

dupE :: Monad m => Enumeratee [a] [(a,a)] m x
dupE = mapChunks (map (\a -> (a,a)))

type I = Iteratee

-- split/combine are meant to work as a pair, and combine probably doesn't
-- behave as expected if it's used on an arbitrary iteratee
--
-- in particular, the iteratees produced by @split@ are meant to receive a chunk
-- at a time, then can be terminated.  The continuation is embedded in the
-- returned iteratee.  @combine@ takes advantage of this in the 'step'
-- function by calling 'run' after feeding data to both monadic iteratee
-- levels.
--
-- With arbitrary iteratees, this may not be what's desired, e.g.
--
-- > *Main Control.Monad> let xs = [(fromIntegral i, i) | i <- [1..10]] :: [(Double,Int)]
-- > *Main Control.Monad> let i = liftM2 (,) stream2list (lift stream2list)
-- > *Main Control.Monad> let i' = combine i
-- > *Main Control.Monad> print =<< run =<< enumPureNChunk xs 3 i'
-- > ([1,2,3,4,5,6,7,8,9,10],[])
-- > *Main Control.Monad> 
--
-- Notice that the inner stream (2nd returned value) is terminated immediately.
combine :: forall a b m x. Monad m => I [b] (I [a] m) x -> I [(a,b)] m x
combine = joinIM . combine'
combine'
      :: forall a b m x. Monad m
      => I [b] (I [a] m) x
      -> m (I [(a,b)] m x)
combine' iter = do
    checkRes <- isDone iter
    case checkRes of
      Just result -> return result
      Nothing     -> return $ liftI (step iter)
  where
    isDone i = runIter (runIter i (\x bStream -> return $ Just (x,bStream))
                                  (\_ _ -> return Nothing))
                       (\res aStream -> case res of
                          Just (x,bStream) -> return . Just $ idone x (zipS aStream bStream)
                          Nothing          -> return Nothing)
                       (\_ _ -> return Nothing)

    zipS (EOF mErr) _ = EOF mErr
    zipS _ (EOF mErr) = EOF mErr
    zipS (Chunk as) (Chunk bs) = Chunk (Prelude.zip as bs)

    step :: I [b] (I [a] m) x -> Stream [(a,b)] -> I [(a,b)] m x
    step i (Chunk tups)
      | null tups = liftI (step i)
      | otherwise = do
          let (as,bs) = unzip tups
          i' <- lift $ (enumPure1Chunk as $ enumPure1Chunk bs i) >>= run
          combine i'
    step i (EOF Nothing) =
      combine =<< (lift $ (enumEof (enumEof i)) >>= run)
    step i (EOF (Just e)) =
      combine =<< (lift $ (enumEof (enumEof i)) >>= run)


-- I know split is good (at least mostly good, it works for all this.  Haven't
-- tested exceptions etc.)
-- > *Main Data.Iteratee Data.Char Data.Function> let i = split (joinI $ takeUpTo 11 stream2list)
-- > *Main Data.Iteratee Data.Char Data.Function> let e1 = enumPureNChunk (cycle "abcde") 2 i
-- > *Main Data.Iteratee Data.Char Data.Function> let e2 = enumPureNChunk (cycle [1..10::Int]) 3 e1
-- >
-- > *Main Data.Iteratee Data.Char Data.Function> :t e2
-- > e2
-- >   :: Monad m =>
-- >      m (Iteratee [Int] m (Iteratee [Char] (I [Int] m) [(Int, Char)]))
-- >
-- > *Main Data.Iteratee Data.Char Data.Function> e2 >>= run >>= (run . run) >>= print
-- > [(1,'a'),(2,'b'),(3,'c'),(4,'d'),(5,'e'),(6,'a'),(7,'b'),(8,'c'),(9,'d'),(10,'e'),(1,'a')]
-- >
-- > *Main Data.Iteratee Data.Char Data.Function> 
--
split :: forall a b m x. (Monad m) => I [(a,b)] m x -> I [b] (I [a] m) x
split = joinIM . lift . split' ([],[])

split' :: forall a b m x. (Monad m) => ([a], [b]) -> I [(a,b)] m x -> m (I [b] (I [a] m) x)
split' (preAs, preBs) iter = do
  checkRes <- isDone iter
  case checkRes of
    Just result -> return result
    Nothing     -> return $ stepper iter
  where
    isDone i = runIter i (\x tups -> return (Just $ finish x tups)) (\k _ -> return Nothing)

    stepper i = do
      ((as,bs), aExc, bExc) <- getChunkLayers preAs preBs
      let (str,rest) = zipRem as bs
      case (aExc,bExc) of
          (Nothing, Nothing) -> joinIM . lift $ enumPure1Chunk str i >>= split' (toAcc rest)
          (Just (EOF e), _)  -> joinIM . lift $ enumPure1Chunk str i >>= enumChunk (EOF e) >>= split' (toAcc rest)
          (_, Just (EOF e))  -> joinIM . lift $ enumPure1Chunk str i >>= enumChunk (EOF e) >>= split' (toAcc rest)
          _ -> error "split: internal error, getChunkLayers did the wrong thing"

    finish :: x -> Stream [(a,b)] -> I [b] (I [a] m) x
    finish x str = let (as,bs) = unTup str in lift (idone () as) >> idone x bs

    unTup (Chunk tups) = let (as,bs) = unzip tups in (Chunk as, Chunk bs)
    unTup (EOF mErr)   = (EOF mErr, EOF mErr)

getChunkLayers :: (Monad m) => [a] -> [b] -> I [b] (I [a] m) (([a],[b]), Maybe (Stream [a]), Maybe (Stream [b]))
getChunkLayers preAs preBs = do
  inner <- lift chunkOrErr
  outer <- chunkOrErr2
  case (inner, outer) of
    (Chunk as, Chunk bs) -> return ((preAs++as,preBs++bs), Nothing, Nothing)
    (a@EOF{}, Chunk bs)  -> return ((preAs,preBs++bs), Just a, Nothing)
    (Chunk as, b@EOF{})  -> return ((preAs++as,preBs),Nothing, Just b)
    (a@EOF{}, b@EOF{})   -> return ((preAs, preBs),Just a, Just b)

-- a forcing getChunk-like function
chunkOrErr2 :: (ListLike s el, NullPoint s, Nullable s, Monad m) => Iteratee s m (Stream s)
chunkOrErr2 = do
  mSz <- chunkLength
  case mSz of
    Nothing -> do
      Just e <- isStreamFinished
      case fromException e of
          Just EofException -> return (EOF Nothing)
          Nothing           -> return (EOF $ Just e)
    Just 0 -> return (Chunk empty)
    Just _ -> Chunk `liftM` getChunk

-- a non-forcing, getChunk-like function
-- returns a stream with data or EOF msg
-- returns Nothing if the current chunk is empty (without forcing it)
chunkOrErr = liftI check
  where
    check s@(Chunk c) = idone s s
    check s@EOF{}     = idone s s

zipRem :: [a] -> [b] -> ([(a,b)], Maybe (Either [a] [b]))
zipRem [] [] = ([], Nothing)
zipRem as [] = ([], Just (Left as))
zipRem [] bs = ([], Just (Right bs))
zipRem (a:as) (b:bs) = let (tl, remain) = zipRem as bs
                       in ((a,b):tl, remain)

toAcc :: Maybe (Either [a] [b]) -> ([a], [b])
toAcc Nothing = ([],[])
toAcc (Just es) = either (,[]) ([],) es