{-# LANGUAGE BangPatterns #-}

module System.IO.Streams.Concurrent.Unagi
       ( -- * Channel conversions
           inputToChan
       , chanToInput
       , chanToOutput
       , concurrentMerge
       , makeChanPipe
       ) where


------------------------------------------------------------------------------
import           Control.Applicative           ((<$>), (<*>))
import           Control.Concurrent            (forkIO)
import           Control.Concurrent.Chan.Unagi (InChan, OutChan, newChan,
                                                readChan, writeChan)
import           Control.Concurrent.MVar       (modifyMVar, newEmptyMVar,
                                                newMVar, putMVar, takeMVar)
import           Control.Exception             (SomeException, mask, throwIO,
                                                try)
import           Control.Monad                 (forM_)
import           Prelude                       hiding (read)
import           System.IO.Streams.Internal    (InputStream, OutputStream,
                                                makeInputStream,
                                                makeOutputStream, read)



------------------------------------------------------------------------------
-- | Writes the contents of an input stream to a channel until the input stream
-- yields end-of-stream.
inputToChan :: InputStream a -> InChan (Maybe a) -> IO ()
inputToChan is ch = go
  where
    go = do
        mb <- read is
        writeChan ch mb
        maybe (return $! ()) (const go) mb


------------------------------------------------------------------------------
-- | Turns an 'OutChan' into an input stream.
--
chanToInput :: OutChan (Maybe a) -> IO (InputStream a)
chanToInput ch = makeInputStream $! readChan ch


------------------------------------------------------------------------------
-- | Turns an 'InChan' into an output stream.
--
chanToOutput :: InChan (Maybe a) -> IO (OutputStream a)
chanToOutput = makeOutputStream . writeChan


------------------------------------------------------------------------------
-- | Concurrently merges a list of 'InputStream's, combining values in the
-- order they become available.
--
-- Note: does /not/ forward individual end-of-stream notifications, the
-- produced stream does not yield end-of-stream until all of the input streams
-- have finished.
--
-- This traps exceptions in each concurrent thread and re-raises them in the
-- current thread.
concurrentMerge :: [InputStream a] -> IO (InputStream a)
concurrentMerge iss = do
    mv    <- newEmptyMVar
    nleft <- newMVar $! length iss
    mask $ \restore -> forM_ iss $ \is -> forkIO $ do
        let producer = do
              emb <- try $ restore $ read is
              case emb of
                  Left exc      -> do putMVar mv (Left (exc :: SomeException))
                                      producer
                  Right Nothing -> putMVar mv $! Right Nothing
                  Right x       -> putMVar mv (Right x) >> producer
        producer
    makeInputStream $ chunk mv nleft

  where
    chunk mv nleft = do
        emb <- takeMVar mv
        case emb of
            Left exc      -> throwIO exc
            Right Nothing -> do x <- modifyMVar nleft $ \n ->
                                     let !n' = n - 1
                                     in return $! (n', n')
                                if x > 0
                                  then chunk mv nleft
                                  else return Nothing
            Right x       -> return x


--------------------------------------------------------------------------------
-- | Create a new pair of streams using an underlying 'Chan'. Everything written
-- to the 'OutputStream' will appear as-is on the 'InputStream'.
--
-- Since reading from the 'InputStream' and writing to the 'OutputStream' are
-- blocking calls, be sure to do so in different threads.
makeChanPipe :: IO (InputStream a, OutputStream a)
makeChanPipe = do
    (inChan, outChan) <- newChan
    (,) <$> chanToInput outChan <*> chanToOutput inChan