{-# 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 x <- InputStream a -> IO (Maybe a)
forall a. InputStream a -> IO (Maybe a)
Streams.read InputStream a
i
          Streams.write x o
          return 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 ref <- Text -> IO (IORef Text)
forall a. a -> IO (IORef a)
newIORef Text
forall a. Monoid a => a
mempty
       Streams.makeOutputStream (con 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 start <- IORef Text -> IO Text
forall a. IORef a -> IO a
readIORef IORef Text
ref
      case 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 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
     out_str <- Streams.decodeUtf8 =<< Streams.handleToInputStream out_h
     err_reader <- startHandleReader err_h Nothing
     return (in_str, out_str, err_reader)
demuxProcessHandles Handle
in_h Handle
out_h Handle
err_h (Just (Text
comment_prefix, Handle
aux_h)) =
  do 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
     in_str  <- Streams.encodeUtf8 =<< Streams.handleToOutputStream in_h
     out_str <- Streams.decodeUtf8 =<< Streams.handleToInputStream out_h

     in_aux <- lineBufferedOutputStream mempty aux_str
     in_str' <- teeOutputStream in_str in_aux

     out_aux <- lineBufferedOutputStream comment_prefix aux_str
     out_str' <- teeInputStream out_str out_aux

     err_reader <- startHandleReader err_h . Just
                    =<< lineBufferedOutputStream comment_prefix aux_str

     return (in_str', out_str', 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 ln <- Handle -> IO Text
Text.hGetLine Handle
h
         writeChan c (Just ln)
         go
streamLines Chan (Maybe Text)
c Handle
h (Just OutputStream Text
auxstr) = IO ()
forall {b}. IO b
go
 where
 go :: IO b
go = do ln <- Handle -> IO Text
Text.hGetLine Handle
h
         Streams.write (Just ln) auxstr
         writeChan c (Just ln)
         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
  c <- IO (Chan (Maybe Text))
forall a. IO (Chan a)
newChan
  let 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
  tid <- forkIO $ streamLines c h auxOutput `catch` handle_err

  return $! HandleReader { hrChan     = c
                         , hrHandle   = h
                         , hrThreadId = 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
  mr <- Chan (Maybe Text) -> IO (Maybe Text)
forall a. Chan a -> IO a
readChan (HandleReader -> Chan (Maybe Text)
hrChan HandleReader
hr)
  case 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()
  return 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
          mr <- HandleReader -> IO (Maybe Text)
readNextLine HandleReader
hr
          case 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'