{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Module: Capnp.IO
-- Description: Utilities for reading and writing values to handles.
--
-- This module provides utilities for reading and writing values to and
-- from file 'Handle's.
module Capnp.IO
  ( sGetMsg,
    sPutMsg,
    M.hGetMsg,
    M.getMsg,
    M.hPutMsg,
    M.putMsg,
    hGetParsed,
    sGetParsed,
    getParsed,
    hPutParsed,
    sPutParsed,
    putParsed,
    hGetRaw,
    getRaw,
    sGetRaw,
  )
where

import Capnp.Bits (WordCount, wordsToBytes)
import Capnp.Classes (Parse)
import Capnp.Convert
  ( msgToLBS,
    msgToParsed,
    msgToRaw,
    parsedToBuilder,
    parsedToLBS,
  )
import Capnp.Message (Mutability (..))
import qualified Capnp.Message as M
import qualified Capnp.Repr as R
import Capnp.TraversalLimit (evalLimitT)
import Control.Exception (throwIO)
import Control.Monad.Trans.Class (lift)
import Data.Bits
import qualified Data.ByteString as BS
import qualified Data.ByteString.Builder as BB
import Network.Simple.TCP (Socket, recv, sendLazy)
import System.IO (Handle, stdin, stdout)
import System.IO.Error (eofErrorType, mkIOError)

-- | Like 'hGetMsg', except that it takes a socket instead of a 'Handle'.
sGetMsg :: Socket -> WordCount -> IO (M.Message 'Const)
sGetMsg :: Socket -> WordCount -> IO (Message 'Const)
sGetMsg Socket
socket WordCount
limit =
  forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT WordCount
limit forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
(MonadThrow m, MonadLimit m) =>
m Word32 -> (WordCount -> m (Segment 'Const)) -> m (Message 'Const)
M.readMessage (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift IO Word32
read32) (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. WordCount -> IO (Segment 'Const)
readSegment)
  where
    read32 :: IO Word32
read32 = do
      ByteString
bytes <- Int -> IO ByteString
recvFull Int
4
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
        (forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString
bytes HasCallStack => ByteString -> Int -> Word8
`BS.index` Int
0) forall a. Bits a => a -> Int -> a
`shiftL` Int
0)
          forall a. Bits a => a -> a -> a
.|. (forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString
bytes HasCallStack => ByteString -> Int -> Word8
`BS.index` Int
1) forall a. Bits a => a -> Int -> a
`shiftL` Int
8)
          forall a. Bits a => a -> a -> a
.|. (forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString
bytes HasCallStack => ByteString -> Int -> Word8
`BS.index` Int
2) forall a. Bits a => a -> Int -> a
`shiftL` Int
16)
          forall a. Bits a => a -> a -> a
.|. (forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString
bytes HasCallStack => ByteString -> Int -> Word8
`BS.index` Int
3) forall a. Bits a => a -> Int -> a
`shiftL` Int
24)
    readSegment :: WordCount -> IO (Segment 'Const)
readSegment !WordCount
words =
      ByteString -> Segment 'Const
M.fromByteString forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO ByteString
recvFull (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ WordCount -> ByteCount
wordsToBytes WordCount
words)

    -- \| Like recv, but (1) never returns less than `count` bytes, (2)
    -- uses `socket`, rather than taking the socket as an argument, and (3)
    -- throws an EOF exception when the connection is closed.
    recvFull :: Int -> IO BS.ByteString
    recvFull :: Int -> IO ByteString
recvFull !Int
count = do
      Maybe ByteString
maybeBytes <- forall (m :: * -> *).
MonadIO m =>
Socket -> Int -> m (Maybe ByteString)
recv Socket
socket Int
count
      case Maybe ByteString
maybeBytes of
        Maybe ByteString
Nothing ->
          forall e a. Exception e => e -> IO a
throwIO forall a b. (a -> b) -> a -> b
$ IOErrorType -> String -> Maybe Handle -> Maybe String -> IOError
mkIOError IOErrorType
eofErrorType String
"Remote socket closed" forall a. Maybe a
Nothing forall a. Maybe a
Nothing
        Just ByteString
bytes
          | ByteString -> Int
BS.length ByteString
bytes forall a. Eq a => a -> a -> Bool
== Int
count ->
              forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
bytes
          | Bool
otherwise ->
              (ByteString
bytes forall a. Semigroup a => a -> a -> a
<>) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO ByteString
recvFull (Int
count forall a. Num a => a -> a -> a
- ByteString -> Int
BS.length ByteString
bytes)

-- | Like 'hPutMsg', except that it takes a 'Socket' instead of a 'Handle'.
sPutMsg :: Socket -> M.Message 'Const -> IO ()
sPutMsg :: Socket -> Message 'Const -> IO ()
sPutMsg Socket
socket = forall (m :: * -> *). MonadIO m => Socket -> ByteString -> m ()
sendLazy Socket
socket forall b c a. (b -> c) -> (a -> b) -> a -> c
. Message 'Const -> ByteString
msgToLBS

-- | Read a struct from the handle in its parsed form, using the supplied
-- read limit.
hGetParsed :: forall a pa. (R.IsStruct a, Parse a pa) => Handle -> WordCount -> IO pa
hGetParsed :: forall a pa.
(IsStruct a, Parse a pa) =>
Handle -> WordCount -> IO pa
hGetParsed Handle
handle WordCount
limit = do
  Message 'Const
msg <- Handle -> WordCount -> IO (Message 'Const)
M.hGetMsg Handle
handle WordCount
limit
  forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT WordCount
limit forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *) pa.
(ReadCtx m 'Const, IsStruct a, Parse a pa) =>
Message 'Const -> m pa
msgToParsed @a Message 'Const
msg

-- | Read a struct from the socket in its parsed form, using the supplied
-- read limit.
sGetParsed :: forall a pa. (R.IsStruct a, Parse a pa) => Socket -> WordCount -> IO pa
sGetParsed :: forall a pa.
(IsStruct a, Parse a pa) =>
Socket -> WordCount -> IO pa
sGetParsed Socket
socket WordCount
limit = do
  Message 'Const
msg <- Socket -> WordCount -> IO (Message 'Const)
sGetMsg Socket
socket WordCount
limit
  forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT WordCount
limit forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *) pa.
(ReadCtx m 'Const, IsStruct a, Parse a pa) =>
Message 'Const -> m pa
msgToParsed @a Message 'Const
msg

-- | Read a struct from stdin in its parsed form, using the supplied
-- read limit.
getParsed :: (R.IsStruct a, Parse a pa) => WordCount -> IO pa
getParsed :: forall a pa. (IsStruct a, Parse a pa) => WordCount -> IO pa
getParsed = forall a pa.
(IsStruct a, Parse a pa) =>
Handle -> WordCount -> IO pa
hGetParsed Handle
stdin

-- | Write the parsed form of a struct to the handle
hPutParsed :: (R.IsStruct a, Parse a pa) => Handle -> pa -> IO ()
hPutParsed :: forall a pa. (IsStruct a, Parse a pa) => Handle -> pa -> IO ()
hPutParsed Handle
h pa
value = do
  Builder
bb <- forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT forall a. Bounded a => a
maxBound forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *) pa s.
(RWCtx m s, IsStruct a, Parse a pa) =>
pa -> m Builder
parsedToBuilder pa
value
  Handle -> Builder -> IO ()
BB.hPutBuilder Handle
h Builder
bb

-- | Write the parsed form of a struct to stdout
putParsed :: (R.IsStruct a, Parse a pa) => pa -> IO ()
putParsed :: forall a pa. (IsStruct a, Parse a pa) => pa -> IO ()
putParsed = forall a pa. (IsStruct a, Parse a pa) => Handle -> pa -> IO ()
hPutParsed Handle
stdout

-- | Write the parsed form of a struct to the socket.
sPutParsed :: (R.IsStruct a, Parse a pa) => Socket -> pa -> IO ()
sPutParsed :: forall a pa. (IsStruct a, Parse a pa) => Socket -> pa -> IO ()
sPutParsed Socket
socket pa
value = do
  ByteString
lbs <- forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT forall a. Bounded a => a
maxBound forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *) pa s.
(RWCtx m s, IsStruct a, Parse a pa) =>
pa -> m ByteString
parsedToLBS pa
value
  forall (m :: * -> *). MonadIO m => Socket -> ByteString -> m ()
sendLazy Socket
socket ByteString
lbs

-- | Read a struct from the handle using the supplied read limit,
-- and return its root pointer.
hGetRaw :: R.IsStruct a => Handle -> WordCount -> IO (R.Raw a 'Const)
hGetRaw :: forall a. IsStruct a => Handle -> WordCount -> IO (Raw a 'Const)
hGetRaw Handle
h WordCount
limit = do
  Message 'Const
msg <- Handle -> WordCount -> IO (Message 'Const)
M.hGetMsg Handle
h WordCount
limit
  forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT WordCount
limit forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *) (mut :: Mutability).
(ReadCtx m mut, IsStruct a) =>
Message mut -> m (Raw a mut)
msgToRaw Message 'Const
msg

-- | Read a struct from stdin using the supplied read limit,
-- and return its root pointer.
getRaw :: R.IsStruct a => WordCount -> IO (R.Raw a 'Const)
getRaw :: forall a. IsStruct a => WordCount -> IO (Raw a 'Const)
getRaw = forall a. IsStruct a => Handle -> WordCount -> IO (Raw a 'Const)
hGetRaw Handle
stdin

-- | Read a struct from the socket using the supplied read limit,
-- and return its root pointer.
sGetRaw :: R.IsStruct a => Socket -> WordCount -> IO (R.Raw a 'Const)
sGetRaw :: forall a. IsStruct a => Socket -> WordCount -> IO (Raw a 'Const)
sGetRaw Socket
socket WordCount
limit = do
  Message 'Const
msg <- Socket -> WordCount -> IO (Message 'Const)
sGetMsg Socket
socket WordCount
limit
  forall (m :: * -> *) a.
MonadThrow m =>
WordCount -> LimitT m a -> m a
evalLimitT WordCount
limit forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *) (mut :: Mutability).
(ReadCtx m mut, IsStruct a) =>
Message mut -> m (Raw a mut)
msgToRaw Message 'Const
msg