{-# 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 :: 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 (m :: Type -> Type) a. Monad m => a -> m a
return Maybe a
x

teeOutputStream :: OutputStream a -> OutputStream a -> IO (OutputStream a)
teeOutputStream :: 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 (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 (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 (m :: Type -> Type) a. Monad m => a -> m a
return (HandleReader -> IO HandleReader)
-> HandleReader -> IO HandleReader
forall a b. (a -> b) -> a -> b
$! HandleReader :: Chan (Maybe Text) -> Handle -> ThreadId -> HandleReader
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 :: 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 (m :: Type -> Type) a. Monad m => a -> m a
return()
  Maybe Text -> IO (Maybe Text)
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 (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'