Synchronized.hs

module Lattest.Streams.Synchronized (
TInputStream,
read,
readAll,
unRead,
hasInput,
fromInputStreamBuffered,
makeTInputStream,
fromBuffer,
fromTMVar,
duplicate,
map,
mapUnbuffered,
mergeBuffered,
mergeBufferedWith,
consumeBufferedWith,
Streamed(..),
tryRead,
tryReadIO,
tryReadIO'
)
where

import Prelude hiding (read, map)
import Control.Concurrent.STM(STM, retry)
import Control.Concurrent.STM.TQueue(TQueue, newTQueueIO, writeTQueue, readTQueue, isEmptyTQueue, unGetTQueue)
import Control.Concurrent.STM.TMVar(TMVar, takeTMVar, isEmptyTMVar)
import Control.Concurrent.STM.TVar(newTVarIO, readTVar, writeTVar)
import GHC.IO.Exception (ioe_location, ioe_errno)
import Control.Exception(throwIO, Exception, catchJust)
import Control.Monad (void, when)
import Data.List(singleton)

import GHC.Conc (atomically, forkIO)
import System.IO.Streams (InputStream, OutputStream, makeOutputStream, connect)
import qualified System.IO.Streams as Streams (write)
import Control.Monad.Extra((||^))
import Foreign.C.Error (eBADF, Errno (Errno))

data TInputStream a = TInputStream {
    tRead :: STM (Maybe a),
    tUnRead :: a -> STM (),
    tHasInput :: STM Bool -- whether this stream has a (Maybe a) available right now, *or* whether it is closed. Note: it may be tempting
                          -- to omit this field, and just check whether tRead retries, but that is ambiguous: tRead will retry either if
                          -- there is no input yet, *or* if someone is writing to the unknown TVars used to implement tRead. The first case
                          -- would not have an input, the second case would.
    }

read :: TInputStream a -> STM (Maybe a)
read = tRead

hasInput :: TInputStream a -> STM Bool
hasInput = tHasInput

-- read inputs from the stream until it has no more input
readAll :: TInputStream a -> STM [Maybe a]
readAll tis = reverse <$> readAll' []
    where
    readAll' acc = do
        readAttempt <- tryRead tis
        case readAttempt of
            StreamClosed -> return (Nothing:acc)
            Unavailable -> return acc
            Available next -> readAll' (Just next:acc)

unRead :: a -> TInputStream a -> STM ()
unRead = flip tUnRead

data Streamed a = Available a | Unavailable | StreamClosed deriving (Eq, Ord, Show, Read)

tryRead :: TInputStream a -> STM (Streamed a)
tryRead tis = do
    ready <- hasInput tis
    if ready
        then do
            minput <- read tis
            return $ maybe StreamClosed Available minput
        else return Unavailable

tryReadIO :: TInputStream a -> IO (Streamed a)
tryReadIO = atomically . tryRead

data TInputStreamClosedException = TInputStreamClosedException deriving (Eq, Ord, Show)
instance Exception TInputStreamClosedException

tryReadIO' :: TInputStream a -> IO (Maybe a)
tryReadIO' tis = do
    streamed <- tryReadIO tis
    case streamed of
        Available a -> return $ Just a
        Unavailable -> return Nothing
        StreamClosed -> throwIO TInputStreamClosedException

makeTInputStream :: STM (Maybe a) -> STM Bool -> IO (TInputStream a)
makeTInputStream readInput hasInput' = do
    pushbackStack <- newTQueueIO
    withClosedState $ TInputStream {
        tRead = do
            hasPushback <- isNonEmptyTQueue pushbackStack
            if hasPushback
                then Just <$> readTQueue pushbackStack
                else readInput,
        tUnRead = unGetTQueue pushbackStack,
        tHasInput = isNonEmptyTQueue pushbackStack ||^ hasInput'
        }

withClosedState :: TInputStream a -> IO (TInputStream a)
withClosedState stream = do
    closeVar <- newTVarIO False
    return $ TInputStream {
        tRead = do
            closed <- readTVar closeVar
            if closed
                then return Nothing
                else do
                    mi <- tRead stream
                    case mi of
                        Just i -> return $ Just i
                        Nothing -> do
                            writeTVar closeVar True
                            return Nothing,
        tUnRead = tUnRead stream,
        tHasInput = do
            closed <- readTVar closeVar
            if closed
                then return True
                else tHasInput stream
        }

fromTMVar :: TMVar (Maybe a) -> IO (TInputStream a)
fromTMVar tmvar = makeTInputStream (takeTMVar tmvar) (not <$> isEmptyTMVar tmvar)

-- note: the buffer is also used to push back values. Thus, do not read this buffer from outside the resulting stream, as then this stream may
-- violate the principle that pushed back values are read again.
fromBuffer :: TQueue (Maybe a) -> IO (TInputStream a)
fromBuffer buffer = withClosedState $ TInputStream {
        tRead = readTQueue buffer,
        tUnRead = unGetTQueue buffer . Just,
        tHasInput = isNonEmptyTQueue buffer
    }

-- a synchronized input stream that reads from the given input stream. A monitoring thread reads the given input stream and buffers the inputs on the background
-- TODO support a bound for the buffer
fromInputStreamBuffered :: InputStream a -> IO (TInputStream a)
fromInputStreamBuffered is = do
    buffer <- newTQueueIO
    bufferOS <- makeOutputStream $ atomically . writeTQueue buffer -- an intermediate output stream to write to the queue
    void $ forkIO $ catchSocketClosed $ connect is bufferOS -- move items from the original adapter into the buffer in a separate thread
    fromBuffer buffer
 where
   -- catches "threadWait: invalid argument (Bad file descriptor)"
   -- which is thrown when the socket is closed while this thread is running, e.g. on Adaptor.close()
   -- The alternative solution is to keep track of the threadID given by forkIO,
   -- and explicitly stop the thread before closing the socket.
   catchSocketClosed forkedThread = catchJust threadwaitinvalid forkedThread pure
   threadwaitinvalid e
    | ioe_location e == "threadWait"
    , fmap Errno (ioe_errno e) == Just eBADF = Just ()
    | otherwise = Nothing

mapUnbuffered :: (a -> b) -> (b -> a) -> TInputStream a -> TInputStream b -- we need an inverse to push a's back to the original TInputStream of b's
mapUnbuffered f f' tis = TInputStream { -- FIXME either remove or wrap in withClosedState
    tRead = fmap f <$> read tis,
    tUnRead = tUnRead tis . f',
    tHasInput = tHasInput tis
    }

map :: (a -> b) -> TInputStream a -> IO (TInputStream b) 
map f tis = makeTInputStream (fmap f <$> read tis) (hasInput tis)

-- Take inputs from both given input streams. If both streams have an input available, the first stream takes precedence. The resulting stream closes 
-- if either given stream closes. The merge is buffered, so that the sequence of arrival of inputs in both streams is preserved, even if the resulting
-- stream is not read.
mergeBuffered :: TInputStream a -> TInputStream a -> IO (TInputStream a)
mergeBuffered = mergeBufferedWith $ \tis1 tis2 -> singleton <$> do
    ready1 <- hasInput tis1
    if ready1
        then read tis1
        else do
            ready2 <- hasInput tis2
            if ready2
                then read tis2
                else retry

-- Take inputs from both given input streams, using the given combinator function. The merge is buffered, so that the sequence of arrival of inputs in both streams is preserved,
-- even if the resulting stream is not read. The resulting stream closes if either given stream closes.
mergeBufferedWith :: (TInputStream a -> TInputStream a -> STM [Maybe a]) -> TInputStream a -> TInputStream a -> IO (TInputStream a)
mergeBufferedWith merge tis1 tis2 = consumeBufferedWith (merge tis1 tis2) (hasAnyInput tis1 tis2) -- FIXME if one stream closes, then no further inputs should be read from the other stream, and only Nothings should be returned instead

hasAnyInput :: TInputStream a -> TInputStream a -> STM Bool
hasAnyInput tis1 tis2 = hasInput tis1 ||^ hasInput tis2

-- Actively take inputs from the given list STM, via a separate thread. This ensures that ordering from the given list STM is preserved as much as possible,
-- in case of concurrent inputs. The resulting stream only reads via the separate thread if needed, and reads from the list STM directly if this is faster.
-- The resulting stream closes if the consumer closes. This function assumes that the given STMs close properly, i.e., no Just's are sent after a Nothing
consumeBufferedWith :: STM [Maybe a] -> STM Bool -> IO (TInputStream a)
consumeBufferedWith producer producerHasInput = do
    buffer <- newTQueueIO
    let writeToBuffer = writeTQueue buffer
    --bufferOS <- makeOutputStream $ \x -> putStrLn "Synchronized.merging" >> (atomically . writeTQueue buffer) x -- an intermediate output stream to write to the queue
    void $ forkIO $ bufferLoop writeToBuffer -- move items from the original streams into the buffer in a separate thread
    bufferIS <- fromBuffer buffer
    return $ bufferIS {
        tRead = pickOneAndMergeRest buffer writeToBuffer,
        tHasInput = hasInput bufferIS ||^ producerHasInput
        }
    where
    bufferLoop writeToBuffer = do
        continue <- atomically $ do
            produced <- producer
            writeSeqTo produced writeToBuffer
        when continue $ bufferLoop writeToBuffer
    writeSeqTo [] _ = return True  -- TODO this is an OutputStream utils function, similar to Streams.write/writeTo, move to an appropriate module
    writeSeqTo (Nothing:_) writeToBuffer = do
        _ <- writeToBuffer Nothing
        return False
    writeSeqTo (Just a:as) writeToBuffer = do
        _ <- writeToBuffer (Just a)
        writeSeqTo as writeToBuffer
    pickOneAndMergeRest buffer writeToBuffer = do
        bufferHasInput <- not <$> isEmptyTQueue buffer
        produced <- if bufferHasInput
            then singleton <$> readTQueue buffer
            else producer
        case produced of
            [] -> retry
            (Just x:xs) -> do
                void $ writeSeqTo xs writeToBuffer
                return $ Just x
            (Nothing:_) -> do
                unGetTQueue buffer Nothing
                return Nothing

{-
-- TODO either remove or expose as nice util function, even though we don't seem to need this now
tconnect :: TInputStream a -> OutputStream a -> IO ()
tconnect tis os = tconnect'
    where
    tconnect' = do
        mA <- atomically $ read tis
        case mA of
            Nothing -> Streams.write Nothing os
            Just _ -> do
                Streams.write mA os
                tconnect'
-}

-- FIXME split to a separate module, conceptually there's nothing synchronized about this w.r.t. other OutputStreams
duplicate :: OutputStream a -> OutputStream a -> IO (OutputStream a)
duplicate out1 out2 = makeOutputStream $ \maybeVal -> do
    Streams.write maybeVal out1
    Streams.write maybeVal out2

isNonEmptyTQueue :: TQueue a -> STM Bool
isNonEmptyTQueue queue = not <$> isEmptyTQueue queue