Forking an enumerator computation

EnumeratorFork.hs

{-# LANGUAGE ScopedTypeVariables #-}

import           Control.Concurrent (killThread)
import           Control.Concurrent.BoundedChan
import           Control.Concurrent.Thread
import           Control.Exception (SomeException)
import           Control.Monad.CatchIO
import           Control.Monad.Trans
import           Prelude hiding (catch)
import qualified Data.Enumerator.List           as E
import           Snap.Iteratee


chanToIter :: BoundedChan (Maybe a) -> Iteratee a IO ()
chanToIter chan = go
  where
    go = E.head >>= 
         maybe (liftIO $ writeChan chan Nothing)
               (\x -> liftIO (writeChan chan (Just x)) >> go)


chanToEnum :: BoundedChan (Maybe a) -> Enumerator a IO b
chanToEnum chan = check
  where
    readUntilEOF = do
        mbX <- liftIO $ readChan chan
        maybe (return ()) (const readUntilEOF) mbX

    getStreamChunk = do
        mbX <- liftIO $ readChan chan
        return $ maybe EOF (Chunks . (:[])) mbX

    check (Continue k) = do
        stream <- getStreamChunk
        step <- lift $ runIteratee $ k stream
        check step

    check (Yield x r) = readUntilEOF >> yield x r

    check (Error e) = throwError e


forkEnumerator :: MonadIO m =>
                  (Enumerator a IO b -> IO c)
               -> m (Iteratee a IO c)
forkEnumerator func = do
    chan <- liftIO $ newBoundedChan chanSize

    (tid, resultThunk) <- liftIO $ forkIO $ func $ chanToEnum chan
    return $ mkOutputIter tid resultThunk chan

  where
    chanSize = 4

    mkOutputIter tid resultThunk chan = do
        res <- outputIter `catch` \(e::SomeException) -> do
                   liftIO $ killThread tid
                   throwError e
        either throwError return res

      where
        outputIter = do
            chanToIter chan
            liftIO resultThunk