{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE BangPatterns #-}


{-|
Module      : Z.IO.StdStream
Description : TTY devices
Copyright   : (c) Dong Han, 2018~2019
License     : BSD
Maintainer  : winterland1989@gmail.com
Stability   : experimental
Portability : non-portable

This module provides stdin\/stderr\/stdout reading and writings. Usually you don't have to use 'stderr' or 'stderrBuf' directly, 'Z.IO.Logger' provides more logging utilities through @stderr@. While 'stdinBuf' and 'stdoutBuf' is useful when you write interactive programs, 'Z.IO.Buffered' module provide many reading and writing operations. Example:

@
import Control.Concurrent.MVar
import Z.IO.LowResTimer
import Z.IO.Buffered
import Z.IO.StdStream
import qualified Z.Data.Vector as V
import qualified Z.Data.Builder as B
main = do
    -- read by '\n'
    b1 <- readLineStd
    -- read whatever user input in 3s, otherwise get Nothing
    b2 <- timeoutLowRes 30 $ withMVar stdinBuf readBuffer
    ...
    putStd "hello world!"

    -- Raw mode
    setStdinTTYMode UV_TTY_MODE_RAW
    forever $ do
        withMVar stdinBuf $ \ i -> withMVar stdoutBuf $ \ o -> do
            bs <- readBuffer i
            let Just key = V.headMaybe bs
            writeBuilder o (B.hex key)
            flushBuffer o
@

-}
module Z.IO.StdStream
  ( -- * Standard input & output streams
    StdStream
  , isStdStreamTTY
  , UVTTYMode(UV_TTY_MODE_NORMAL, UV_TTY_MODE_RAW)
  , setStdinTTYMode
  , getStdoutWinSize
  , stdin, stdout, stderr
  , stdinBuf, stdoutBuf, stderrBuf
    -- * utils
  , printStd
  , readLineStd
  , putStd
  , putLineStd
  ) where

import Control.Monad
import Control.Concurrent.MVar
import Foreign.C.Types (CInt)
import Foreign.Ptr
import System.IO.Unsafe
import Z.Data.Builder as B
import Z.Data.Vector as V
import Z.Data.Text.Builder (ToText, toBuilder)
import Z.IO.UV.FFI
import Z.IO.UV.Manager
import Z.IO.UV.Errno
import Z.IO.Exception
import Z.IO.Buffered
import Z.Foreign

-- | Standard input and output streams
--
-- We support both regular file and TTY based streams, when initialized
-- 'uv_guess_handle' is called to decide which type of devices are connected
-- to standard streams.
--
-- Note 'StdStream' is not thread safe, you shouldn't use them without lock.
-- For the same reason you shouldn't use stderr directly, use `Z.IO.Logger` module instead.

data StdStream
    = StdTTY {-# UNPACK #-}!(Ptr UVHandle) {-# UNPACK #-}!UVSlot UVManager -- similar to UVStream
    | StdFile {-# UNPACK #-}!UVFD                                          -- similar to UVFile

isStdStreamTTY :: StdStream -> Bool
isStdStreamTTY :: StdStream -> Bool
isStdStreamTTY (StdTTY Ptr UVHandle
_ UVSlot
_ UVManager
_) = Bool
True
isStdStreamTTY StdStream
_              = Bool
False

instance Input StdStream where
    {-# INLINE readInput #-}
    readInput :: StdStream -> Ptr Word8 -> UVSlot -> IO UVSlot
readInput (StdTTY Ptr UVHandle
hdl UVSlot
slot UVManager
uvm) Ptr Word8
buf UVSlot
len = IO UVSlot -> IO UVSlot
forall a. IO a -> IO a
mask_ (IO UVSlot -> IO UVSlot) -> IO UVSlot -> IO UVSlot
forall a b. (a -> b) -> a -> b
$ do
        UVManager -> UVSlot -> Ptr Word8 -> UVSlot -> IO ()
pokeBufferTable UVManager
uvm UVSlot
slot Ptr Word8
buf UVSlot
len
        MVar UVSlot
m <- UVManager -> UVSlot -> IO (MVar UVSlot)
getBlockMVar UVManager
uvm UVSlot
slot
        Maybe UVSlot
_ <- MVar UVSlot -> IO (Maybe UVSlot)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar UVSlot
m
        IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ UVManager -> IO CInt -> IO CInt
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm (Ptr UVHandle -> IO CInt
hs_uv_read_start Ptr UVHandle
hdl)
        -- since we are inside mask, this is the only place
        -- async exceptions could possibly kick in, and we should stop reading
        UVSlot
r <- MVar UVSlot -> IO UVSlot
forall a. MVar a -> IO a
takeMVar MVar UVSlot
m IO UVSlot -> IO () -> IO UVSlot
forall a b. IO a -> IO b -> IO a
`onException` (do
                IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ UVManager -> IO CInt -> IO CInt
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm (Ptr UVHandle -> IO CInt
uv_read_stop Ptr UVHandle
hdl)
                IO (Maybe UVSlot) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (MVar UVSlot -> IO (Maybe UVSlot)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar UVSlot
m))
        if  | UVSlot
r UVSlot -> UVSlot -> Bool
forall a. Ord a => a -> a -> Bool
> UVSlot
0  -> UVSlot -> IO UVSlot
forall (m :: * -> *) a. Monad m => a -> m a
return UVSlot
r
            -- r == 0 should be impossible, since we guard this situation in c side
            | UVSlot
r UVSlot -> UVSlot -> Bool
forall a. Eq a => a -> a -> Bool
== CInt -> UVSlot
forall a b. (Integral a, Num b) => a -> b
fromIntegral CInt
UV_EOF -> UVSlot -> IO UVSlot
forall (m :: * -> *) a. Monad m => a -> m a
return UVSlot
0
            | UVSlot
r UVSlot -> UVSlot -> Bool
forall a. Ord a => a -> a -> Bool
< UVSlot
0 ->  IO UVSlot -> IO UVSlot
forall a. (HasCallStack, Integral a) => IO a -> IO a
throwUVIfMinus (UVSlot -> IO UVSlot
forall (m :: * -> *) a. Monad m => a -> m a
return UVSlot
r)
    readInput (StdFile UVFD
fd) Ptr Word8
buf UVSlot
len =
        IO UVSlot -> IO UVSlot
forall a. (HasCallStack, Integral a) => IO a -> IO a
throwUVIfMinus (IO UVSlot -> IO UVSlot) -> IO UVSlot -> IO UVSlot
forall a b. (a -> b) -> a -> b
$ UVFD -> Ptr Word8 -> UVSlot -> Int64 -> IO UVSlot
hs_uv_fs_read UVFD
fd Ptr Word8
buf UVSlot
len (-Int64
1)

instance Output StdStream where
    {-# INLINE writeOutput #-}
    writeOutput :: StdStream -> Ptr Word8 -> UVSlot -> IO ()
writeOutput (StdTTY Ptr UVHandle
hdl UVSlot
_ UVManager
uvm) Ptr Word8
buf UVSlot
len = IO () -> IO ()
forall a. IO a -> IO a
mask_ (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        MVar UVSlot
m <- UVManager -> IO (MVar UVSlot) -> IO (MVar UVSlot)
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm (IO (MVar UVSlot) -> IO (MVar UVSlot))
-> IO (MVar UVSlot) -> IO (MVar UVSlot)
forall a b. (a -> b) -> a -> b
$ do
            UVSlot
reqSlot <- HasCallStack => UVManager -> IO UVSlotUnSafe -> IO UVSlot
UVManager -> IO UVSlotUnSafe -> IO UVSlot
getUVSlot UVManager
uvm (Ptr UVHandle -> Ptr Word8 -> UVSlot -> IO UVSlotUnSafe
hs_uv_write Ptr UVHandle
hdl Ptr Word8
buf UVSlot
len)
            MVar UVSlot
m <- UVManager -> UVSlot -> IO (MVar UVSlot)
getBlockMVar UVManager
uvm UVSlot
reqSlot
            Maybe UVSlot
_ <- MVar UVSlot -> IO (Maybe UVSlot)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar MVar UVSlot
m
            MVar UVSlot -> IO (MVar UVSlot)
forall (m :: * -> *) a. Monad m => a -> m a
return MVar UVSlot
m
        IO UVSlot -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO UVSlot -> IO UVSlot
forall a. IO a -> IO a
uninterruptibleMask_ (IO UVSlot -> IO UVSlot) -> IO UVSlot -> IO UVSlot
forall a b. (a -> b) -> a -> b
$ MVar UVSlot -> IO UVSlot
forall a. MVar a -> IO a
takeMVar MVar UVSlot
m)
    writeOutput (StdFile UVFD
fd) Ptr Word8
buf UVSlot
len = Ptr Word8 -> UVSlot -> IO ()
go Ptr Word8
buf UVSlot
len
      where
        go :: Ptr Word8 -> UVSlot -> IO ()
go !Ptr Word8
b !UVSlot
bufSiz = do
            UVSlot
written <- IO UVSlot -> IO UVSlot
forall a. (HasCallStack, Integral a) => IO a -> IO a
throwUVIfMinus
                (UVFD -> Ptr Word8 -> UVSlot -> Int64 -> IO UVSlot
hs_uv_fs_write UVFD
fd Ptr Word8
b UVSlot
bufSiz (-Int64
1))
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (UVSlot
written UVSlot -> UVSlot -> Bool
forall a. Ord a => a -> a -> Bool
< UVSlot
bufSiz)
                (Ptr Word8 -> UVSlot -> IO ()
go (Ptr Word8
b Ptr Word8 -> UVSlot -> Ptr Word8
forall a b. Ptr a -> UVSlot -> Ptr b
`plusPtr` UVSlot
written) (UVSlot
bufSizUVSlot -> UVSlot -> UVSlot
forall a. Num a => a -> a -> a
-UVSlot
written))

-- | The global stdin stream.
stdin :: StdStream
{-# NOINLINE stdin #-}
stdin :: StdStream
stdin = IO StdStream -> StdStream
forall a. IO a -> a
unsafePerformIO (UVFD -> IO StdStream
makeStdStream UVFD
0)

-- | The global stdout stream.
--
-- | If you want to write logs, don't use 'stdout' directly, use 'Z.IO.Logger' instead.
stdout :: StdStream
{-# NOINLINE stdout #-}
stdout :: StdStream
stdout = IO StdStream -> StdStream
forall a. IO a -> a
unsafePerformIO (UVFD -> IO StdStream
makeStdStream UVFD
1)

-- | The global stderr stream.
--
-- | If you want to write logs, don't use 'stderr' directly, use 'Z.IO.Logger' instead.
stderr :: StdStream
{-# NOINLINE stderr #-}
stderr :: StdStream
stderr = IO StdStream -> StdStream
forall a. IO a -> a
unsafePerformIO (UVFD -> IO StdStream
makeStdStream UVFD
2)

-- |  A global buffered stdin stream protected by 'MVar'.
stdinBuf :: MVar (BufferedInput StdStream)
{-# NOINLINE stdinBuf #-}
stdinBuf :: MVar (BufferedInput StdStream)
stdinBuf = IO (MVar (BufferedInput StdStream))
-> MVar (BufferedInput StdStream)
forall a. IO a -> a
unsafePerformIO (UVSlot -> StdStream -> IO (BufferedInput StdStream)
forall input. UVSlot -> input -> IO (BufferedInput input)
newBufferedInput UVSlot
defaultChunkSize StdStream
stdin IO (BufferedInput StdStream)
-> (BufferedInput StdStream -> IO (MVar (BufferedInput StdStream)))
-> IO (MVar (BufferedInput StdStream))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= BufferedInput StdStream -> IO (MVar (BufferedInput StdStream))
forall a. a -> IO (MVar a)
newMVar)

-- |  A global buffered stdout stream protected by 'MVar'.
--
-- | If you want to write logs, don't use 'stdoutBuf' directly, use 'Z.IO.Logger' instead.
stdoutBuf :: MVar (BufferedOutput StdStream)
{-# NOINLINE stdoutBuf #-}
stdoutBuf :: MVar (BufferedOutput StdStream)
stdoutBuf = IO (MVar (BufferedOutput StdStream))
-> MVar (BufferedOutput StdStream)
forall a. IO a -> a
unsafePerformIO (UVSlot -> StdStream -> IO (BufferedOutput StdStream)
forall output. UVSlot -> output -> IO (BufferedOutput output)
newBufferedOutput UVSlot
defaultChunkSize StdStream
stdout IO (BufferedOutput StdStream)
-> (BufferedOutput StdStream
    -> IO (MVar (BufferedOutput StdStream)))
-> IO (MVar (BufferedOutput StdStream))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= BufferedOutput StdStream -> IO (MVar (BufferedOutput StdStream))
forall a. a -> IO (MVar a)
newMVar)

-- |  A global buffered stderr stream protected by 'MVar'.
--
-- | If you want to write logs, don't use 'stderrBuf' directly, use 'Z.IO.Logger' instead.
stderrBuf :: MVar (BufferedOutput StdStream)
{-# NOINLINE stderrBuf #-}
stderrBuf :: MVar (BufferedOutput StdStream)
stderrBuf = IO (MVar (BufferedOutput StdStream))
-> MVar (BufferedOutput StdStream)
forall a. IO a -> a
unsafePerformIO (UVSlot -> StdStream -> IO (BufferedOutput StdStream)
forall output. UVSlot -> output -> IO (BufferedOutput output)
newBufferedOutput UVSlot
defaultChunkSize StdStream
stderr IO (BufferedOutput StdStream)
-> (BufferedOutput StdStream
    -> IO (MVar (BufferedOutput StdStream)))
-> IO (MVar (BufferedOutput StdStream))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= BufferedOutput StdStream -> IO (MVar (BufferedOutput StdStream))
forall a. a -> IO (MVar a)
newMVar)

makeStdStream :: UVFD -> IO StdStream
makeStdStream :: UVFD -> IO StdStream
makeStdStream UVFD
fd = do
    UVHandleType
typ <- UVFD -> IO UVHandleType
uv_guess_handle UVFD
fd
    if UVHandleType
typ UVHandleType -> UVHandleType -> Bool
forall a. Eq a => a -> a -> Bool
== UVHandleType
UV_TTY
    then do
        UVManager
uvm <- IO UVManager
getUVManager
        UVManager -> (Ptr UVLoop -> IO StdStream) -> IO StdStream
forall a. HasCallStack => UVManager -> (Ptr UVLoop -> IO a) -> IO a
withUVManager UVManager
uvm ((Ptr UVLoop -> IO StdStream) -> IO StdStream)
-> (Ptr UVLoop -> IO StdStream) -> IO StdStream
forall a b. (a -> b) -> a -> b
$ \ Ptr UVLoop
loop -> do
            Ptr UVHandle
hdl <- Ptr UVLoop -> IO (Ptr UVHandle)
hs_uv_handle_alloc Ptr UVLoop
loop
            UVSlot
slot <- HasCallStack => UVManager -> IO UVSlotUnSafe -> IO UVSlot
UVManager -> IO UVSlotUnSafe -> IO UVSlot
getUVSlot UVManager
uvm (Ptr UVHandle -> IO UVSlotUnSafe
peekUVHandleData Ptr UVHandle
hdl)
            Maybe UVSlot
_ <- MVar UVSlot -> IO (Maybe UVSlot)
forall a. MVar a -> IO (Maybe a)
tryTakeMVar (MVar UVSlot -> IO (Maybe UVSlot))
-> IO (MVar UVSlot) -> IO (Maybe UVSlot)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< UVManager -> UVSlot -> IO (MVar UVSlot)
getBlockMVar UVManager
uvm UVSlot
slot   -- clear the parking spot
            IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (Ptr UVLoop -> Ptr UVHandle -> CInt -> IO CInt
uv_tty_init Ptr UVLoop
loop Ptr UVHandle
hdl (UVFD -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral UVFD
fd))
                IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO a
`onException` Ptr UVHandle -> IO ()
hs_uv_handle_free Ptr UVHandle
hdl
            StdStream -> IO StdStream
forall (m :: * -> *) a. Monad m => a -> m a
return (Ptr UVHandle -> UVSlot -> UVManager -> StdStream
StdTTY Ptr UVHandle
hdl UVSlot
slot UVManager
uvm)
    else StdStream -> IO StdStream
forall (m :: * -> *) a. Monad m => a -> m a
return (UVFD -> StdStream
StdFile UVFD
fd)

-- | Change terminal's mode if stdin is connected to a terminal.
setStdinTTYMode :: UVTTYMode -> IO ()
setStdinTTYMode :: UVTTYMode -> IO ()
setStdinTTYMode UVTTYMode
mode = case StdStream
stdin of
    StdTTY Ptr UVHandle
hdl UVSlot
_ UVManager
uvm ->
        UVManager -> IO () -> IO ()
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm (IO () -> IO ()) -> (IO CInt -> IO ()) -> IO CInt -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr UVHandle -> UVTTYMode -> IO CInt
uv_tty_set_mode Ptr UVHandle
hdl UVTTYMode
mode
    StdStream
_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Get terminal's output window size in (width, height) format,
-- return (-1, -1) if stdout is a file.
getStdoutWinSize :: IO (CInt, CInt)
getStdoutWinSize :: IO (CInt, CInt)
getStdoutWinSize = case StdStream
stdout of
    StdTTY Ptr UVHandle
hdl UVSlot
_ UVManager
uvm ->
        UVManager -> IO (CInt, CInt) -> IO (CInt, CInt)
forall a. HasCallStack => UVManager -> IO a -> IO a
withUVManager' UVManager
uvm (IO (CInt, CInt) -> IO (CInt, CInt))
-> IO (CInt, CInt) -> IO (CInt, CInt)
forall a b. (a -> b) -> a -> b
$ do
            (CInt
w, (CInt
h, ())) <- (MBA# CInt -> IO (CInt, ())) -> IO (CInt, (CInt, ()))
forall a b. Prim a => (MBA# CInt -> IO b) -> IO (a, b)
allocPrimUnsafe ((MBA# CInt -> IO (CInt, ())) -> IO (CInt, (CInt, ())))
-> (MBA# CInt -> IO (CInt, ())) -> IO (CInt, (CInt, ()))
forall a b. (a -> b) -> a -> b
$ \ MBA# CInt
w ->
                (MBA# CInt -> IO ()) -> IO (CInt, ())
forall a b. Prim a => (MBA# CInt -> IO b) -> IO (a, b)
allocPrimUnsafe ((MBA# CInt -> IO ()) -> IO (CInt, ()))
-> (MBA# CInt -> IO ()) -> IO (CInt, ())
forall a b. (a -> b) -> a -> b
$ \ MBA# CInt
h -> IO CInt -> IO ()
forall a. (HasCallStack, Integral a) => IO a -> IO ()
throwUVIfMinus_ (IO CInt -> IO ()) -> IO CInt -> IO ()
forall a b. (a -> b) -> a -> b
$ Ptr UVHandle -> MBA# CInt -> MBA# CInt -> IO CInt
uv_tty_get_winsize Ptr UVHandle
hdl MBA# CInt
w MBA# CInt
h
            (CInt, CInt) -> IO (CInt, CInt)
forall (m :: * -> *) a. Monad m => a -> m a
return (CInt
w, CInt
h)
    StdStream
_ -> (CInt, CInt) -> IO (CInt, CInt)
forall (m :: * -> *) a. Monad m => a -> m a
return (-CInt
1, -CInt
1)

--------------------------------------------------------------------------------

-- | print a 'ToText' and flush to stdout.
printStd :: ToText a => a -> IO ()
printStd :: a -> IO ()
printStd a
s = Builder () -> IO ()
forall a. Builder a -> IO ()
putStd (a -> Builder ()
forall a. ToText a => a -> Builder ()
toBuilder a
s)

-- | print a 'Builder' and flush to stdout.
putStd :: Builder a -> IO ()
putStd :: Builder a -> IO ()
putStd Builder a
b = MVar (BufferedOutput StdStream)
-> (BufferedOutput StdStream -> IO ()) -> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (BufferedOutput StdStream)
stdoutBuf ((BufferedOutput StdStream -> IO ()) -> IO ())
-> (BufferedOutput StdStream -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ BufferedOutput StdStream
o -> do
    BufferedOutput StdStream -> Builder a -> IO ()
forall o a. Output o => BufferedOutput o -> Builder a -> IO ()
writeBuilder BufferedOutput StdStream
o Builder a
b
    BufferedOutput StdStream -> IO ()
forall f. Output f => BufferedOutput f -> IO ()
flushBuffer BufferedOutput StdStream
o

-- | print a 'Builder' and flush to stdout, with a linefeed.
putLineStd :: Builder a -> IO ()
putLineStd :: Builder a -> IO ()
putLineStd Builder a
b = MVar (BufferedOutput StdStream)
-> (BufferedOutput StdStream -> IO ()) -> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (BufferedOutput StdStream)
stdoutBuf ((BufferedOutput StdStream -> IO ()) -> IO ())
-> (BufferedOutput StdStream -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ BufferedOutput StdStream
o -> do
    BufferedOutput StdStream -> Builder () -> IO ()
forall o a. Output o => BufferedOutput o -> Builder a -> IO ()
writeBuilder BufferedOutput StdStream
o (Builder a
b Builder a -> Builder () -> Builder ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Char -> Builder ()
B.char8 Char
'\n')
    BufferedOutput StdStream -> IO ()
forall f. Output f => BufferedOutput f -> IO ()
flushBuffer BufferedOutput StdStream
o

-- | read a line from stdin
readLineStd :: IO V.Bytes
readLineStd :: IO Bytes
readLineStd = MVar (BufferedInput StdStream)
-> (BufferedInput StdStream -> IO Bytes) -> IO Bytes
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar (BufferedInput StdStream)
stdinBuf BufferedInput StdStream -> IO Bytes
forall i. (HasCallStack, Input i) => BufferedInput i -> IO Bytes
readLine