    {-# LANGUAGE ScopedTypeVariables #-}
    module What4.Utils.HandleReader where

    import           Control.Monad (unless)
    import           Data.IORef
    import           Data.Text (Text)
    import qualified Data.Text as Text
    import qualified Data.Text.Lazy as LazyText
    import qualified Data.Text.IO as Text
    import           Control.Exception(bracket,catch,IOException)
    import           Control.Concurrent(ThreadId,forkIO,killThread)
    import           Control.Concurrent.Chan(Chan,newChan,readChan,writeChan)
    import           System.IO(Handle,hClose)
    import           System.IO.Streams( OutputStream, InputStream )
    import qualified System.IO.Streams as Streams

    teeInputStream :: InputStream a -> OutputStream a -> IO (InputStream a)
    teeInputStream :: forall a. InputStream a -> OutputStream a -> IO (InputStream a)
    teeInputStream InputStream a
    i OutputStream a
    o = IO (Maybe a) -> IO (InputStream a)
    forall a. IO (Maybe a) -> IO (InputStream a)
    Streams.makeInputStream IO (Maybe a)
    go
      where
      go :: IO (Maybe a)
    go = do Maybe a
    x <- InputStream a -> IO (Maybe a)
    forall a. InputStream a -> IO (Maybe a)
    Streams.read InputStream a
    i
              Maybe a -> OutputStream a -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write Maybe a
    x OutputStream a
    o
              Maybe a -> IO (Maybe a)
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return Maybe a
    x

    teeOutputStream :: OutputStream a -> OutputStream a -> IO (OutputStream a)
    teeOutputStream :: forall a. OutputStream a -> OutputStream a -> IO (OutputStream a)
    teeOutputStream OutputStream a
    o OutputStream a
    aux = (Maybe a -> IO ()) -> IO (OutputStream a)
    forall a. (Maybe a -> IO ()) -> IO (OutputStream a)
    Streams.makeOutputStream Maybe a -> IO ()
    go
      where
      go :: Maybe a -> IO ()
    go Maybe a
    x =
        do Maybe a -> OutputStream a -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write Maybe a
    x OutputStream a
    aux
           Maybe a -> OutputStream a -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write Maybe a
    x OutputStream a
    o

    lineBufferedOutputStream :: Text -> OutputStream Text -> IO (OutputStream Text)
    lineBufferedOutputStream :: Text -> OutputStream Text -> IO (OutputStream Text)
    lineBufferedOutputStream Text
    prefix OutputStream Text
    out =
        do IORef Text
    ref <- Text -> IO (IORef Text)
    forall a. a -> IO (IORef a)
    newIORef Text
    forall a. Monoid a => a
    mempty
           (Maybe Text -> IO ()) -> IO (OutputStream Text)
    forall a. (Maybe a -> IO ()) -> IO (OutputStream a)
    Streams.makeOutputStream (IORef Text -> Maybe Text -> IO ()
    con IORef Text
    ref)
     where
     newl :: Text
    newl = String -> Text
    Text.pack String
    "\n"

     con :: IORef Text -> Maybe Text -> IO ()
    con IORef Text
    ref Maybe Text
    mx =
       do Text
    start <- IORef Text -> IO Text
    forall a. IORef a -> IO a
    readIORef IORef Text
    ref
          case Maybe Text
    mx of
            Maybe Text
    Nothing ->
              do Bool -> IO () -> IO ()
    forall (f :: Type -> Type). Applicative f => Bool -> f () -> f ()
    unless (Text -> Bool
    Text.null Text
    start) (Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write (Text -> Maybe Text
    forall a. a -> Maybe a
    Just (Text
    prefix Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    start)) OutputStream Text
    out)
                 Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write Maybe Text
    forall a. Maybe a
    Nothing OutputStream Text
    out
            Just Text
    x -> IORef Text -> Text -> IO ()
    go IORef Text
    ref (Text
    start Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    x)

     go :: IORef Text -> Text -> IO ()
    go IORef Text
    ref Text
    x =
       let (Text
    ln, Text
    x') = (Char -> Bool) -> Text -> (Text, Text)
    Text.break (Char -> Char -> Bool
    forall a. Eq a => a -> a -> Bool
    == Char
    '\n') Text
    x in
       if Text -> Bool
    Text.null Text
    x' then
         -- Flush
         do Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write (Text -> Maybe Text
    forall a. a -> Maybe a
    Just Text
    forall a. Monoid a => a
    mempty) OutputStream Text
    out
            IORef Text -> Text -> IO ()
    forall a. IORef a -> a -> IO ()
    writeIORef IORef Text
    ref Text
    x
       else
         do Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write (Text -> Maybe Text
    forall a. a -> Maybe a
    Just (Text
    prefix Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    ln Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    newl)) OutputStream Text
    out
            IORef Text -> Text -> IO ()
    go IORef Text
    ref (Int -> Text -> Text
    Text.drop Int
    1 Text
    x')

    demuxProcessHandles ::
      Handle {- ^ stdin for process -} ->
      Handle {- ^ stdout for process -} ->
      Handle {- ^ stderr for process -} ->
      Maybe (Text, Handle) {- optional handle to echo ouput; text argument is a line-comment prefix  -} ->
      IO ( OutputStream Text, InputStream Text, HandleReader )
    demuxProcessHandles :: Handle
    -> Handle
    -> Handle
    -> Maybe (Text, Handle)
    -> IO (OutputStream Text, InputStream Text, HandleReader)
    demuxProcessHandles Handle
    in_h Handle
    out_h Handle
    err_h Maybe (Text, Handle)
    Nothing =
      do OutputStream Text
    in_str  <- OutputStream ByteString -> IO (OutputStream Text)
    Streams.encodeUtf8 (OutputStream ByteString -> IO (OutputStream Text))
    -> IO (OutputStream ByteString) -> IO (OutputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Handle -> IO (OutputStream ByteString)
    Streams.handleToOutputStream Handle
    in_h
         InputStream Text
    out_str <- InputStream ByteString -> IO (InputStream Text)
    Streams.decodeUtf8 (InputStream ByteString -> IO (InputStream Text))
    -> IO (InputStream ByteString) -> IO (InputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Handle -> IO (InputStream ByteString)
    Streams.handleToInputStream Handle
    out_h
         HandleReader
    err_reader <- Handle -> Maybe (OutputStream Text) -> IO HandleReader
    startHandleReader Handle
    err_h Maybe (OutputStream Text)
    forall a. Maybe a
    Nothing
         (OutputStream Text, InputStream Text, HandleReader)
    -> IO (OutputStream Text, InputStream Text, HandleReader)
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return (OutputStream Text
    in_str, InputStream Text
    out_str, HandleReader
    err_reader)
    demuxProcessHandles Handle
    in_h Handle
    out_h Handle
    err_h (Just (Text
    comment_prefix, Handle
    aux_h)) =
      do OutputStream Text
    aux_str <- OutputStream Text -> IO (OutputStream Text)
    forall a. OutputStream a -> IO (OutputStream a)
    Streams.lockingOutputStream (OutputStream Text -> IO (OutputStream Text))
    -> IO (OutputStream Text) -> IO (OutputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< OutputStream ByteString -> IO (OutputStream Text)
    Streams.encodeUtf8 (OutputStream ByteString -> IO (OutputStream Text))
    -> IO (OutputStream ByteString) -> IO (OutputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Handle -> IO (OutputStream ByteString)
    Streams.handleToOutputStream Handle
    aux_h
         OutputStream Text
    in_str  <- OutputStream ByteString -> IO (OutputStream Text)
    Streams.encodeUtf8 (OutputStream ByteString -> IO (OutputStream Text))
    -> IO (OutputStream ByteString) -> IO (OutputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Handle -> IO (OutputStream ByteString)
    Streams.handleToOutputStream Handle
    in_h
         InputStream Text
    out_str <- InputStream ByteString -> IO (InputStream Text)
    Streams.decodeUtf8 (InputStream ByteString -> IO (InputStream Text))
    -> IO (InputStream ByteString) -> IO (InputStream Text)
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Handle -> IO (InputStream ByteString)
    Streams.handleToInputStream Handle
    out_h

         OutputStream Text
    in_aux <- Text -> OutputStream Text -> IO (OutputStream Text)
    lineBufferedOutputStream Text
    forall a. Monoid a => a
    mempty OutputStream Text
    aux_str
         OutputStream Text
    in_str' <- OutputStream Text -> OutputStream Text -> IO (OutputStream Text)
    forall a. OutputStream a -> OutputStream a -> IO (OutputStream a)
    teeOutputStream OutputStream Text
    in_str OutputStream Text
    in_aux

         OutputStream Text
    out_aux <- Text -> OutputStream Text -> IO (OutputStream Text)
    lineBufferedOutputStream Text
    comment_prefix OutputStream Text
    aux_str
         InputStream Text
    out_str' <- InputStream Text -> OutputStream Text -> IO (InputStream Text)
    forall a. InputStream a -> OutputStream a -> IO (InputStream a)
    teeInputStream InputStream Text
    out_str OutputStream Text
    out_aux

         HandleReader
    err_reader <- Handle -> Maybe (OutputStream Text) -> IO HandleReader
    startHandleReader Handle
    err_h (Maybe (OutputStream Text) -> IO HandleReader)
    -> (OutputStream Text -> Maybe (OutputStream Text))
    -> OutputStream Text
    -> IO HandleReader
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . OutputStream Text -> Maybe (OutputStream Text)
    forall a. a -> Maybe a
    Just
                        (OutputStream Text -> IO HandleReader)
    -> IO (OutputStream Text) -> IO HandleReader
    forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Text -> OutputStream Text -> IO (OutputStream Text)
    lineBufferedOutputStream Text
    comment_prefix OutputStream Text
    aux_str

         (OutputStream Text, InputStream Text, HandleReader)
    -> IO (OutputStream Text, InputStream Text, HandleReader)
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return (OutputStream Text
    in_str', InputStream Text
    out_str', HandleReader
    err_reader)

    {- | Wrapper to help with reading from another process's
         standard out and stderr.

    We want to be able to read from another process's stderr and stdout without
    causing the process to stall because 'stdout' or 'stderr' becomes full.  This
    data type will read from either of the handles, and buffer as much data
    as needed in the queue.  It then provides a line-based method for reading
    that data as strict bytestrings. -}
    data HandleReader = HandleReader { HandleReader -> Chan (Maybe Text)
    hrChan :: !(Chan (Maybe Text))
                                     , HandleReader -> Handle
    hrHandle :: !Handle
                                     , HandleReader -> ThreadId
    hrThreadId :: !ThreadId
                                     }

    streamLines :: Chan (Maybe Text) -> Handle -> Maybe (OutputStream Text) -> IO ()
    streamLines :: Chan (Maybe Text) -> Handle -> Maybe (OutputStream Text) -> IO ()
    streamLines Chan (Maybe Text)
    c Handle
    h Maybe (OutputStream Text)
    Nothing = IO ()
    forall {b}. IO b
    go
     where
     go :: IO b
    go = do Text
    ln <- Handle -> IO Text
    Text.hGetLine Handle
    h
             Chan (Maybe Text) -> Maybe Text -> IO ()
    forall a. Chan a -> a -> IO ()
    writeChan Chan (Maybe Text)
    c (Text -> Maybe Text
    forall a. a -> Maybe a
    Just Text
    ln)
             IO b
    go
    streamLines Chan (Maybe Text)
    c Handle
    h (Just OutputStream Text
    auxstr) = IO ()
    forall {b}. IO b
    go
     where
     go :: IO b
    go = do Text
    ln <- Handle -> IO Text
    Text.hGetLine Handle
    h
             Maybe Text -> OutputStream Text -> IO ()
    forall a. Maybe a -> OutputStream a -> IO ()
    Streams.write (Text -> Maybe Text
    forall a. a -> Maybe a
    Just Text
    ln) OutputStream Text
    auxstr
             Chan (Maybe Text) -> Maybe Text -> IO ()
    forall a. Chan a -> a -> IO ()
    writeChan Chan (Maybe Text)
    c (Text -> Maybe Text
    forall a. a -> Maybe a
    Just Text
    ln)
             IO b
    go

    -- | Create a new handle reader for reading the given handle.
    startHandleReader :: Handle -> Maybe (OutputStream Text) -> IO HandleReader
    startHandleReader :: Handle -> Maybe (OutputStream Text) -> IO HandleReader
    startHandleReader Handle
    h Maybe (OutputStream Text)
    auxOutput = do
      Chan (Maybe Text)
    c <- IO (Chan (Maybe Text))
    forall a. IO (Chan a)
    newChan
      let handle_err :: IOException -> IO ()
    handle_err (IOException
    _e :: IOException) = Chan (Maybe Text) -> Maybe Text -> IO ()
    forall a. Chan a -> a -> IO ()
    writeChan Chan (Maybe Text)
    c Maybe Text
    forall a. Maybe a
    Nothing
      ThreadId
    tid <- IO () -> IO ThreadId
    forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
    forall a b. (a -> b) -> a -> b
    $ Chan (Maybe Text) -> Handle -> Maybe (OutputStream Text) -> IO ()
    streamLines Chan (Maybe Text)
    c Handle
    h Maybe (OutputStream Text)
    auxOutput IO () -> (IOException -> IO ()) -> IO ()
    forall e a. Exception e => IO a -> (e -> IO a) -> IO a
    `catch` IOException -> IO ()
    handle_err

      HandleReader -> IO HandleReader
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return (HandleReader -> IO HandleReader)
    -> HandleReader -> IO HandleReader
    forall a b. (a -> b) -> a -> b
    $! HandleReader { hrChan :: Chan (Maybe Text)
    hrChan     = Chan (Maybe Text)
    c
                             , hrHandle :: Handle
    hrHandle   = Handle
    h
                             , hrThreadId :: ThreadId
    hrThreadId = ThreadId
    tid
                             }

    -- | Stop the handle reader; cannot be used afterwards.
    stopHandleReader :: HandleReader -> IO ()
    stopHandleReader :: HandleReader -> IO ()
    stopHandleReader HandleReader
    hr = do
      ThreadId -> IO ()
    killThread (HandleReader -> ThreadId
    hrThreadId HandleReader
    hr)
      Handle -> IO ()
    hClose (HandleReader -> Handle
    hrHandle HandleReader
    hr)

    -- | Run an execution with a handle reader and stop it wheen down
    withHandleReader :: Handle -> Maybe (OutputStream Text) -> (HandleReader -> IO a) -> IO a
    withHandleReader :: forall a.
    Handle
    -> Maybe (OutputStream Text) -> (HandleReader -> IO a) -> IO a
    withHandleReader Handle
    h Maybe (OutputStream Text)
    auxOut = IO HandleReader
    -> (HandleReader -> IO ()) -> (HandleReader -> IO a) -> IO a
    forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
    bracket (Handle -> Maybe (OutputStream Text) -> IO HandleReader
    startHandleReader Handle
    h Maybe (OutputStream Text)
    auxOut) HandleReader -> IO ()
    stopHandleReader

    readNextLine :: HandleReader -> IO (Maybe Text)
    readNextLine :: HandleReader -> IO (Maybe Text)
    readNextLine HandleReader
    hr = do
      Maybe Text
    mr <- Chan (Maybe Text) -> IO (Maybe Text)
    forall a. Chan a -> IO a
    readChan (HandleReader -> Chan (Maybe Text)
    hrChan HandleReader
    hr)
      case Maybe Text
    mr of
        -- Write back 'Nothing' because thread should have terminated.
        Maybe Text
    Nothing -> Chan (Maybe Text) -> Maybe Text -> IO ()
    forall a. Chan a -> a -> IO ()
    writeChan (HandleReader -> Chan (Maybe Text)
    hrChan HandleReader
    hr) Maybe Text
    forall a. Maybe a
    Nothing
        Just{} -> () -> IO ()
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return()
      Maybe Text -> IO (Maybe Text)
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return Maybe Text
    mr

    readAllLines :: HandleReader -> IO LazyText.Text
    readAllLines :: HandleReader -> IO Text
    readAllLines HandleReader
    hr = Text -> IO Text
    go Text
    LazyText.empty
      where go :: LazyText.Text -> IO LazyText.Text
            go :: Text -> IO Text
    go Text
    prev = do
              Maybe Text
    mr <- HandleReader -> IO (Maybe Text)
    readNextLine HandleReader
    hr
              case Maybe Text
    mr of
                Maybe Text
    Nothing -> Text -> IO Text
    forall a. a -> IO a
    forall (m :: Type -> Type) a. Monad m => a -> m a
    return Text
    prev
                Just Text
    e -> Text -> IO Text
    go (Text -> IO Text) -> Text -> IO Text
    forall a b. (a -> b) -> a -> b
    $! Text
    prev Text -> Text -> Text
    `LazyText.append` (Text -> Text
    LazyText.fromStrict Text
    e)
                                     Text -> Char -> Text
    `LazyText.snoc` Char
    '\n'