{-# LANGUAGE RecordWildCards #-}
module Database.Franz.Internal.IO (getInt64le, runGetRecv, hGetRange) where

import Data.IORef
import Data.Serialize hiding (getInt64le)
import Data.Typeable (cast)
import Foreign.Ptr
import Foreign.ForeignPtr
import Foreign.C.Types
import GHC.IO.Exception
import GHC.IO.Handle.Internals (withHandle_)
import GHC.IO.Handle.Types (Handle__(..), Handle)
import Network.Socket as S
import Network.Socket.ByteString as SB
import System.Endian (fromLE64)
import System.IO.Error
import System.Posix.Types (Fd(..))
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as B
import qualified GHC.IO.FD as FD
import Data.Word (Word8)

-- | Better implementation of 'Data.Serialize.getInt64le'
getInt64le :: Num a => Get a
getInt64le :: Get a
getInt64le = Word64 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> a) -> (Word64 -> Word64) -> Word64 -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Word64
fromLE64 (Word64 -> a) -> Get Word64 -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word64
getWord64host
{-# INLINE getInt64le #-}

runGetRecv :: IORef B.ByteString -> S.Socket -> Get a -> IO (Either String a)
runGetRecv :: IORef ByteString -> Socket -> Get a -> IO (Either String a)
runGetRecv IORef ByteString
refBuf Socket
sock Get a
m = do
  ByteString
lo <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef IORef ByteString
refBuf
  let go :: Result b -> IO (Either String b)
go (Done b
a ByteString
lo') = do
        IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
refBuf ByteString
lo'
        Either String b -> IO (Either String b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String b -> IO (Either String b))
-> Either String b -> IO (Either String b)
forall a b. (a -> b) -> a -> b
$ b -> Either String b
forall a b. b -> Either a b
Right b
a
      go (Partial ByteString -> Result b
cont) = Socket -> Int -> IO ByteString
SB.recv Socket
sock Int
4096 IO ByteString
-> (ByteString -> IO (Either String b)) -> IO (Either String b)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Result b -> IO (Either String b)
go (Result b -> IO (Either String b))
-> (ByteString -> Result b) -> ByteString -> IO (Either String b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Result b
cont
      go (Fail String
str ByteString
lo') = do
        IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
refBuf ByteString
lo'
        Either String b -> IO (Either String b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String b -> IO (Either String b))
-> Either String b -> IO (Either String b)
forall a b. (a -> b) -> a -> b
$ String -> Either String b
forall a b. a -> Either a b
Left (String -> Either String b) -> String -> Either String b
forall a b. (a -> b) -> a -> b
$ Socket -> String
forall a. Show a => a -> String
show Socket
sock String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
str
  ByteString
bs <- if ByteString -> Bool
B.null ByteString
lo
    then Socket -> Int -> IO ByteString
SB.recv Socket
sock Int
4096
    else ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
lo
  Result a -> IO (Either String a)
forall b. Result b -> IO (Either String b)
go (Result a -> IO (Either String a))
-> Result a -> IO (Either String a)
forall a b. (a -> b) -> a -> b
$ Get a -> ByteString -> Result a
forall a. Get a -> ByteString -> Result a
runGetPartial Get a
m ByteString
bs

withFd :: Handle -> (Fd -> IO a) -> IO a
withFd :: Handle -> (Fd -> IO a) -> IO a
withFd Handle
h Fd -> IO a
f = String -> Handle -> (Handle__ -> IO a) -> IO a
forall a. String -> Handle -> (Handle__ -> IO a) -> IO a
withHandle_ String
"withFd" Handle
h ((Handle__ -> IO a) -> IO a) -> (Handle__ -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \ Handle__{dev
Maybe (TextEncoder enc_state)
Maybe (TextDecoder dec_state)
Maybe TextEncoding
Maybe (MVar Handle__)
HandleType
BufferMode
Newline
IORef (dec_state, Buffer Word8)
IORef (BufferList CharBufElem)
IORef (Buffer CharBufElem)
IORef (Buffer Word8)
haDevice :: ()
haType :: Handle__ -> HandleType
haByteBuffer :: Handle__ -> IORef (Buffer Word8)
haBufferMode :: Handle__ -> BufferMode
haLastDecode :: ()
haCharBuffer :: Handle__ -> IORef (Buffer CharBufElem)
haBuffers :: Handle__ -> IORef (BufferList CharBufElem)
haEncoder :: ()
haDecoder :: ()
haCodec :: Handle__ -> Maybe TextEncoding
haInputNL :: Handle__ -> Newline
haOutputNL :: Handle__ -> Newline
haOtherSide :: Handle__ -> Maybe (MVar Handle__)
haOtherSide :: Maybe (MVar Handle__)
haOutputNL :: Newline
haInputNL :: Newline
haCodec :: Maybe TextEncoding
haDecoder :: Maybe (TextDecoder dec_state)
haEncoder :: Maybe (TextEncoder enc_state)
haBuffers :: IORef (BufferList CharBufElem)
haCharBuffer :: IORef (Buffer CharBufElem)
haLastDecode :: IORef (dec_state, Buffer Word8)
haBufferMode :: BufferMode
haByteBuffer :: IORef (Buffer Word8)
haType :: HandleType
haDevice :: dev
..} -> do
  case dev -> Maybe FD
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast dev
haDevice of
    Maybe FD
Nothing -> IOError -> IO a
forall a. IOError -> IO a
ioError (IOError -> String -> IOError
ioeSetErrorString (IOErrorType -> String -> Maybe Handle -> Maybe String -> IOError
mkIOError IOErrorType
IllegalOperation
                                           String
"withFd" (Handle -> Maybe Handle
forall a. a -> Maybe a
Just Handle
h) Maybe String
forall a. Maybe a
Nothing)
                        String
"handle is not a file descriptor")
    Just FD
fd -> Fd -> IO a
f (CInt -> Fd
Fd (CInt -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral (FD -> CInt
FD.fdFD FD
fd)))

foreign import ccall safe "pread"
  c_pread :: Fd -> Ptr Word8 -> CSize -> CSize -> IO CSize

hGetRange :: Handle -> Int -> Int -> IO B.ByteString
hGetRange :: Handle -> Int -> Int -> IO ByteString
hGetRange Handle
h Int
len Int
ofs = do
  ForeignPtr Word8
fptr <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
len
  CSize
count <- Handle -> (Fd -> IO CSize) -> IO CSize
forall a. Handle -> (Fd -> IO a) -> IO a
withFd Handle
h ((Fd -> IO CSize) -> IO CSize) -> (Fd -> IO CSize) -> IO CSize
forall a b. (a -> b) -> a -> b
$ \Fd
fd -> ForeignPtr Word8 -> (Ptr Word8 -> IO CSize) -> IO CSize
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr ((Ptr Word8 -> IO CSize) -> IO CSize)
-> (Ptr Word8 -> IO CSize) -> IO CSize
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> Fd -> Ptr Word8 -> CSize -> CSize -> IO CSize
c_pread Fd
fd Ptr Word8
ptr (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len) (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ofs)
  ByteString -> IO ByteString
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> Int -> ByteString
B.PS ForeignPtr Word8
fptr Int
0 (Int -> ByteString) -> Int -> ByteString
forall a b. (a -> b) -> a -> b
$ CSize -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral CSize
count