{-# LANGUAGE OverloadedStrings #-}

module Network.Socket.BufferPool.Recv (
    receive,
    receiveNoWait,
    receiveWith,
    makeRecvN,
) where

import qualified Data.ByteString as BS
import Data.ByteString.Internal (ByteString (..), unsafeCreate)
import Data.IORef
import Control.Concurrent.STM (STM)
import Network.Socket (Socket, recvBuf, recvBufNoWait, recvBufWith)

import Network.Socket.BufferPool.Buffer
import Network.Socket.BufferPool.Types

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

-- | The receiving function with a buffer pool.
--   The buffer pool is automatically managed.
--
-- /Caveats of 'BufferPool' apply./
receive :: Socket -> BufferPool -> Recv
receive :: Socket -> BufferPool -> Recv
receive Socket
sock BufferPool
pool = BufferPool -> (Buffer -> Int -> IO Int) -> Recv
withBufferPool BufferPool
pool ((Buffer -> Int -> IO Int) -> Recv)
-> (Buffer -> Int -> IO Int) -> Recv
forall a b. (a -> b) -> a -> b
$ \Buffer
ptr Int
size -> Socket -> Buffer -> Int -> IO Int
recvBuf Socket
sock Buffer
ptr Int
size

-- | Like 'receive' but never blocks and never involves the IO manager:
--   'Nothing' means no data was available (or an error occurred, which a
--   subsequent blocking 'receive' will report properly). @Just \"\"@ is EOF.
--
-- /Caveats of 'BufferPool' apply./
receiveNoWait :: Socket -> BufferPool -> IO (Maybe ByteString)
receiveNoWait :: Socket -> BufferPool -> IO (Maybe ByteString)
receiveNoWait Socket
sock BufferPool
pool = BufferPool -> (Buffer -> Int -> IO Int) -> IO (Maybe ByteString)
tryWithBufferPool BufferPool
pool ((Buffer -> Int -> IO Int) -> IO (Maybe ByteString))
-> (Buffer -> Int -> IO Int) -> IO (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ \Buffer
ptr Int
size ->
    -- Both EAGAIN and real errors map to a negative result, deferring
    -- to the blocking path so errors surface there with their errno
    -- intact.
    Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int) -> IO Int -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Socket -> Buffer -> Int -> IO Int
recvBufNoWait Socket
sock Buffer
ptr Int
size

-- | Like 'receive' but the caller can give up on it: whichever settles
--   first decides, 'Left' when the 'STM' action fires, 'Right' when data
--   arrives.  Data already waiting is taken without blocking at all and
--   so wins over a caller which has given up, as a completed operation
--   does in 'Network.Socket.withReadable'; the give-up is read there
--   and nowhere else.
--
--   Nothing is closed and no thread is interrupted to make this return.
--
-- /Caveats of 'BufferPool' apply./
receiveWith :: Socket -> BufferPool -> STM a -> IO (Either a ByteString)
receiveWith :: forall a. Socket -> BufferPool -> STM a -> IO (Either a ByteString)
receiveWith Socket
sock BufferPool
pool STM a
giveUp = do
    Maybe ByteString
early <- Socket -> BufferPool -> IO (Maybe ByteString)
receiveNoWait Socket
sock BufferPool
pool
    case Maybe ByteString
early of
        Just ByteString
bs -> Either a ByteString -> IO (Either a ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either a ByteString -> IO (Either a ByteString))
-> Either a ByteString -> IO (Either a ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString -> Either a ByteString
forall a b. b -> Either a b
Right ByteString
bs
        Maybe ByteString
Nothing -> BufferPool
-> (Buffer -> Int -> IO (Either a Int)) -> IO (Either a ByteString)
forall a.
BufferPool
-> (Buffer -> Int -> IO (Either a Int)) -> IO (Either a ByteString)
eitherWithBufferPool BufferPool
pool ((Buffer -> Int -> IO (Either a Int)) -> IO (Either a ByteString))
-> (Buffer -> Int -> IO (Either a Int)) -> IO (Either a ByteString)
forall a b. (a -> b) -> a -> b
$ \Buffer
ptr Int
size ->
            Socket -> Buffer -> Int -> STM a -> IO (Either a Int)
forall a. Socket -> Buffer -> Int -> STM a -> IO (Either a Int)
recvBufWith Socket
sock Buffer
ptr Int
size STM a
giveUp

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

-- | This function returns a receiving function
--   based on two receiving functions.
--   The returned function receives exactly N bytes.
--   The first argument is an initial received data.
--   After consuming the initial data, the two functions is used.
--   When N is less than equal to 4096, the buffer pool is used.
--   Otherwise, a new buffer is allocated.
--   In this case, the global lock is taken.
--
-- >>> :seti -XOverloadedStrings
-- >>> tryRecvN "a" 3 =<< _iorefRecv ["bcd"]
-- ("abc","d")
-- >>> tryRecvN "a" 3 =<< _iorefRecv ["bc"]
-- ("abc","")
-- >>> tryRecvN "a" 3 =<< _iorefRecv ["b"]
-- ("ab","")
--
-- /The resulting 'RecvN' is designed to be used in one thread only./
makeRecvN :: ByteString -> Recv -> IO RecvN
makeRecvN :: ByteString -> Recv -> IO RecvN
makeRecvN ByteString
bs0 Recv
recv = do
    IORef ByteString
ref <- ByteString -> IO (IORef ByteString)
forall a. a -> IO (IORef a)
newIORef ByteString
bs0
    RecvN -> IO RecvN
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (RecvN -> IO RecvN) -> RecvN -> IO RecvN
forall a b. (a -> b) -> a -> b
$ IORef ByteString -> Recv -> RecvN
recvN IORef ByteString
ref Recv
recv

-- | The receiving function which receives exactly N bytes
--   (the fourth argument).
recvN :: IORef ByteString -> Recv -> RecvN
recvN :: IORef ByteString -> Recv -> RecvN
recvN IORef ByteString
ref Recv
recv Int
size = do
    ByteString
cached <- IORef ByteString -> Recv
forall a. IORef a -> IO a
readIORef IORef ByteString
ref
    (ByteString
bs, ByteString
leftover) <- ByteString -> Int -> Recv -> IO (ByteString, ByteString)
tryRecvN ByteString
cached Int
size Recv
recv
    IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef ByteString
ref ByteString
leftover
    ByteString -> Recv
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs

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

tryRecvN :: ByteString -> Int -> IO ByteString -> IO (ByteString, ByteString)
tryRecvN :: ByteString -> Int -> Recv -> IO (ByteString, ByteString)
tryRecvN ByteString
init0 Int
siz0 Recv
recv
    | Int
siz0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
len0 = (ByteString, ByteString) -> IO (ByteString, ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((ByteString, ByteString) -> IO (ByteString, ByteString))
-> (ByteString, ByteString) -> IO (ByteString, ByteString)
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
siz0 ByteString
init0
    | Bool
otherwise = ([ByteString] -> [ByteString])
-> Int -> IO (ByteString, ByteString)
go (ByteString
init0 ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:) (Int
siz0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len0)
  where
    len0 :: Int
len0 = ByteString -> Int
BS.length ByteString
init0
    go :: ([ByteString] -> [ByteString])
-> Int -> IO (ByteString, ByteString)
go [ByteString] -> [ByteString]
build Int
left = do
        ByteString
bs <- Recv
recv
        let len :: Int
len = ByteString -> Int
BS.length ByteString
bs
        if Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
            then do
                let cs :: ByteString
cs = Int -> [ByteString] -> ByteString
concatN (Int
siz0 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
left) ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
build []
                (ByteString, ByteString) -> IO (ByteString, ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
cs, ByteString
"")
            else
                if Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
left
                    then do
                        let (ByteString
consume, ByteString
leftover) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
left ByteString
bs
                            ret :: ByteString
ret = Int -> [ByteString] -> ByteString
concatN Int
siz0 ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
build [ByteString
consume]
                        (ByteString, ByteString) -> IO (ByteString, ByteString)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
ret, ByteString
leftover)
                    else do
                        let build' :: [ByteString] -> [ByteString]
build' = [ByteString] -> [ByteString]
build ([ByteString] -> [ByteString])
-> ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString
bs ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:)
                            left' :: Int
left' = Int
left Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len
                        ([ByteString] -> [ByteString])
-> Int -> IO (ByteString, ByteString)
go [ByteString] -> [ByteString]
build' Int
left'

concatN :: Int -> [ByteString] -> ByteString
-- Just because it's logical
concatN :: Int -> [ByteString] -> ByteString
concatN Int
_ [] = ByteString
""
-- To avoid a copy if there's only one ByteString
concatN Int
_ [ByteString
bs] = ByteString
bs
concatN Int
total [ByteString]
bss0 =
    Int -> (Buffer -> IO ()) -> ByteString
unsafeCreate Int
total ((Buffer -> IO ()) -> ByteString)
-> (Buffer -> IO ()) -> ByteString
forall a b. (a -> b) -> a -> b
$ \Buffer
ptr -> [ByteString] -> Buffer -> IO ()
goCopy [ByteString]
bss0 Buffer
ptr
  where
    goCopy :: [ByteString] -> Buffer -> IO ()
goCopy [] Buffer
_ = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    goCopy (ByteString
bs : [ByteString]
bss) Buffer
ptr = do
        Buffer
ptr' <- Buffer -> ByteString -> IO Buffer
copy Buffer
ptr ByteString
bs
        [ByteString] -> Buffer -> IO ()
goCopy [ByteString]
bss Buffer
ptr'

-- | doctest only. Elements in the argument must not be empty.
_iorefRecv :: [ByteString] -> IO (IO ByteString)
_iorefRecv :: [ByteString] -> IO Recv
_iorefRecv [ByteString]
ini = do
    IORef [ByteString]
ref <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef [ByteString]
ini
    Recv -> IO Recv
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Recv -> IO Recv) -> Recv -> IO Recv
forall a b. (a -> b) -> a -> b
$ IORef [ByteString] -> Recv
forall {b}. IsString b => IORef [b] -> IO b
recv IORef [ByteString]
ref
  where
    recv :: IORef [b] -> IO b
recv IORef [b]
ref = do
        [b]
xxs <- IORef [b] -> IO [b]
forall a. IORef a -> IO a
readIORef IORef [b]
ref
        case [b]
xxs of
            [] -> do
                IORef [b] -> [b] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef [b]
ref ([b] -> IO ()) -> [b] -> IO ()
forall a b. (a -> b) -> a -> b
$ [Char] -> [b]
forall a. HasCallStack => [Char] -> a
error [Char]
"closed"
                b -> IO b
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return b
""
            b
x : [b]
xs -> do
                IORef [b] -> [b] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef [b]
ref [b]
xs
                b -> IO b
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return b
x