{-# LANGUAGE ScopedTypeVariables #-}
{-@ LIQUID "--no-termination" @-}
{-@ LIQUID "--short-names"    @-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE CPP #-}
{-|
Module: System.IO.ByteBuffer
Description: Provides an efficient buffering abstraction.

A 'ByteBuffer' is a simple buffer for bytes.  It supports two
operations: refilling with the contents of a 'ByteString', and
consuming a fixed number of bytes.

It is implemented as a pointer, together with counters that keep track
of the offset and the number of bytes in the buffer.  Note that the
counters are simple 'IORef's, so 'ByteBuffer's are not thread-safe!

A 'ByteBuffer' is constructed by 'new' with a given starting length,
and will grow (by repeatedly multiplying its size by 1.5) whenever it
is being fed a 'ByteString' that is too large.
-}

module System.IO.ByteBuffer
       ( ByteBuffer
         -- * Allocation and Deallocation
       , new, free, with
         -- * Query for number of available bytes
       , totalSize, isEmpty, availableBytes
         -- * Feeding new input
       , copyByteString
#ifndef mingw32_HOST_OS
       , fillFromFd
#endif
         -- * Consuming bytes from the buffer
       , consume, unsafeConsume
         -- * Exceptions
       , ByteBufferException (..)
       ) where

import           Control.Applicative
import           Control.Exception (SomeException, throwIO)
import           Control.Exception.Lifted (Exception, bracket, catch)
import qualified Control.Monad.Fail as Fail
import           Control.Monad.IO.Class (MonadIO, liftIO)
import           Control.Monad.Trans.Control (MonadBaseControl)
import           Data.ByteString (ByteString)
import qualified Data.ByteString.Internal as BS
import           Data.IORef
import           Data.Maybe (fromMaybe)
import           Data.Typeable (Typeable)
import           Data.Word
import           Foreign.ForeignPtr
import qualified Foreign.Marshal.Alloc as Alloc
import           Foreign.Marshal.Utils (copyBytes, moveBytes)
import           GHC.Ptr
import           Prelude
import qualified Foreign.C.Error as CE
import           Foreign.C.Types
import           System.Posix.Types (Fd (..))

-- | A buffer into which bytes can be written.
--
-- Invariants:
--
-- * @size >= containedBytes >= consumedBytes >= 0@
--
-- * The range from @ptr@ to @ptr `plusPtr` size@ will be allocated
--
-- * The range from @ptr@ to @ptr `plusPtr` containedBytes@ will
--   contain bytes previously copied to the buffer
--
-- * The buffer contains @containedBytes - consumedBytes@ bytes of
--   data that have been copied to it, but not yet read.  They are in
--   the range from @ptr `plusPtr` consumedBytes@ to @ptr `plusPtr`
--   containedBytes@.
--
-- The first two invariants are encoded in Liquid Haskell, and can
-- be statically checked.
--
-- If an Exception occurs during an operation that modifies a
-- 'ByteBuffer', the 'ByteBuffer' is invalidated and can no longer be
-- used.  Trying to access the 'ByteBuffer' subsequently will result
-- in a 'ByteBufferException' being thrown.
{-@
data BBRef = BBRef
    { size :: {v: Int | v >= 0 }
    , contained :: { v: Int | v >= 0 && v <= size }
    , consumed :: { v: Int | v >= 0 && v <= contained }
    , ptr :: { v: Ptr Word8 | (plen v) = size }
    }
@-}

data BBRef = BBRef {
      BBRef -> Int
size      :: {-# UNPACK #-} !Int
      -- ^ The amount of memory allocated.
    , BBRef -> Int
contained :: {-# UNPACK #-} !Int
      -- ^ The number of bytes that the 'ByteBuffer' currently holds.
    , BBRef -> Int
consumed  :: {-# UNPACK #-} !Int
      -- ^ The number of bytes that have already been consumed.
    , BBRef -> Ptr Word8
ptr       :: {-# UNPACK #-} !(Ptr Word8)
      -- ^ This points to the beginning of the memory allocated for
      -- the 'ByteBuffer'
    }

-- | Exception that is thrown when an invalid 'ByteBuffer' is being used that is no longer valid.
--
-- A 'ByteBuffer' is considered to be invalid if
--
-- - it has explicitly been freed
-- - an Exception has occured during an operation that modified it
data ByteBufferException = ByteBufferException
    { ByteBufferException -> String
_bbeLocation :: !String
      -- ^ function name that caused the exception
    , ByteBufferException -> String
_bbeException :: !String
      -- ^ printed representation of the exception
    }
    deriving (Typeable, ByteBufferException -> ByteBufferException -> Bool
(ByteBufferException -> ByteBufferException -> Bool)
-> (ByteBufferException -> ByteBufferException -> Bool)
-> Eq ByteBufferException
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ByteBufferException -> ByteBufferException -> Bool
== :: ByteBufferException -> ByteBufferException -> Bool
$c/= :: ByteBufferException -> ByteBufferException -> Bool
/= :: ByteBufferException -> ByteBufferException -> Bool
Eq)
instance Show ByteBufferException where
    show :: ByteBufferException -> String
show (ByteBufferException String
loc String
e) = [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
        [String
"ByteBufferException: ByteBuffer was invalidated because of Exception thrown in "
        , String
loc , String
": ", String
e]
instance Exception ByteBufferException

type ByteBuffer = IORef (Either ByteBufferException BBRef)

-- | On any Exception, this will invalidate the ByteBuffer and re-throw the Exception.
--
-- Invalidating the 'ByteBuffer' includes freeing the underlying pointer.
bbHandler :: MonadIO m
    => String
       -- ^ location information: function from which the exception was thrown
    -> ByteBuffer
       -- ^ this 'ByteBuffer' will be invalidated when an Exception occurs
    -> (BBRef -> IO a)
    -> m a
bbHandler :: forall (m :: * -> *) a.
MonadIO m =>
String -> ByteBuffer -> (BBRef -> IO a) -> m a
bbHandler String
loc ByteBuffer
bb BBRef -> IO a
f = IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ (BBRef -> IO a) -> ByteBuffer -> IO a
forall a. (BBRef -> IO a) -> ByteBuffer -> IO a
useBBRef BBRef -> IO a
f ByteBuffer
bb IO a -> (SomeException -> IO a) -> IO a
forall (m :: * -> *) e a.
(MonadBaseControl IO m, Exception e) =>
m a -> (e -> m a) -> m a
`catch` \(SomeException
e :: SomeException) -> do
    ByteBuffer -> IO (Either ByteBufferException BBRef)
forall a. IORef a -> IO a
readIORef ByteBuffer
bb IO (Either ByteBufferException BBRef)
-> (Either ByteBufferException BBRef -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Right BBRef
bbref -> do
            Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
Alloc.free (BBRef -> Ptr Word8
ptr BBRef
bbref)
            ByteBuffer -> Either ByteBufferException BBRef -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef ByteBuffer
bb (ByteBufferException -> Either ByteBufferException BBRef
forall a b. a -> Either a b
Left (ByteBufferException -> Either ByteBufferException BBRef)
-> ByteBufferException -> Either ByteBufferException BBRef
forall a b. (a -> b) -> a -> b
$ String -> String -> ByteBufferException
ByteBufferException String
loc (SomeException -> String
forall a. Show a => a -> String
show SomeException
e))
        Left ByteBufferException
_ -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    SomeException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SomeException
e

-- | Try to use the 'BBRef' of a 'ByteBuffer', or throw a 'ByteBufferException' if it's invalid.
useBBRef :: (BBRef -> IO a) -> ByteBuffer -> IO a
useBBRef :: forall a. (BBRef -> IO a) -> ByteBuffer -> IO a
useBBRef BBRef -> IO a
f ByteBuffer
bb = ByteBuffer -> IO (Either ByteBufferException BBRef)
forall a. IORef a -> IO a
readIORef ByteBuffer
bb IO (Either ByteBufferException BBRef)
-> (Either ByteBufferException BBRef -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ByteBufferException -> IO a)
-> (BBRef -> IO a) -> Either ByteBufferException BBRef -> IO a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either ByteBufferException -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO BBRef -> IO a
f
{-# INLINE useBBRef #-}

totalSize :: MonadIO m => ByteBuffer -> m Int
totalSize :: forall (m :: * -> *). MonadIO m => ByteBuffer -> m Int
totalSize = IO Int -> m Int
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> (ByteBuffer -> IO Int) -> ByteBuffer -> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (BBRef -> IO Int) -> ByteBuffer -> IO Int
forall a. (BBRef -> IO a) -> ByteBuffer -> IO a
useBBRef (Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> IO Int) -> (BBRef -> Int) -> BBRef -> IO Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BBRef -> Int
size)
{-# INLINE totalSize #-}

isEmpty :: MonadIO m => ByteBuffer -> m Bool
isEmpty :: forall (m :: * -> *). MonadIO m => ByteBuffer -> m Bool
isEmpty ByteBuffer
bb = IO Bool -> m Bool
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Bool -> m Bool) -> IO Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0) (Int -> Bool) -> IO Int -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ByteBuffer -> IO Int
forall (m :: * -> *). MonadIO m => ByteBuffer -> m Int
availableBytes ByteBuffer
bb
{-# INLINE isEmpty #-}

-- | Number of available bytes in a 'ByteBuffer' (that is, bytes that
-- have been copied to, but not yet read from the 'ByteBuffer'.
{-@ availableBytes :: MonadIO m => ByteBuffer -> m {v: Int | v >= 0} @-}
availableBytes :: MonadIO m => ByteBuffer -> m Int
availableBytes :: forall (m :: * -> *). MonadIO m => ByteBuffer -> m Int
availableBytes = IO Int -> m Int
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Int -> m Int) -> (ByteBuffer -> IO Int) -> ByteBuffer -> m Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (BBRef -> IO Int) -> ByteBuffer -> IO Int
forall a. (BBRef -> IO a) -> ByteBuffer -> IO a
useBBRef (\BBRef{Int
Ptr Word8
size :: BBRef -> Int
contained :: BBRef -> Int
consumed :: BBRef -> Int
ptr :: BBRef -> Ptr Word8
size :: Int
contained :: Int
consumed :: Int
ptr :: Ptr Word8
..} -> Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
contained Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
consumed))
{-# INLINE availableBytes #-}

-- | Allocates a new ByteBuffer with a given buffer size filling from
-- the given FillBuffer.
--
-- Note that 'ByteBuffer's created with 'new' have to be deallocated
-- explicitly using 'free'.  For automatic deallocation, consider
-- using 'with' instead.
new :: MonadIO m
    => Maybe Int
    -- ^ Size of buffer to allocate.  If 'Nothing', use the default
    -- value of 4MB
    -> m ByteBuffer
    -- ^ The byte buffer.
new :: forall (m :: * -> *). MonadIO m => Maybe Int -> m ByteBuffer
new Maybe Int
ml = IO ByteBuffer -> m ByteBuffer
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteBuffer -> m ByteBuffer) -> IO ByteBuffer -> m ByteBuffer
forall a b. (a -> b) -> a -> b
$ do
    let l :: Int
l = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int -> Int) -> (Maybe Int -> Int) -> Maybe Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe (Int
4Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
1024Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
1024) (Maybe Int -> Int) -> Maybe Int -> Int
forall a b. (a -> b) -> a -> b
$ Maybe Int
ml
    newPtr <- Int -> IO (Ptr Word8)
forall a. Int -> IO (Ptr a)
Alloc.mallocBytes Int
l
    newIORef $ Right BBRef
        { ptr = newPtr
        , size = l
        , contained = 0
        , consumed = 0
        }

-- | Free a byte buffer.
free :: MonadIO m => ByteBuffer -> m ()
free :: forall (m :: * -> *). MonadIO m => ByteBuffer -> m ()
free ByteBuffer
bb = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ ByteBuffer -> IO (Either ByteBufferException BBRef)
forall a. IORef a -> IO a
readIORef ByteBuffer
bb IO (Either ByteBufferException BBRef)
-> (Either ByteBufferException BBRef -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Right BBRef
bbref -> do
        Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
Alloc.free (Ptr Word8 -> IO ()) -> Ptr Word8 -> IO ()
forall a b. (a -> b) -> a -> b
$ BBRef -> Ptr Word8
ptr BBRef
bbref
        ByteBuffer -> Either ByteBufferException BBRef -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef ByteBuffer
bb (Either ByteBufferException BBRef -> IO ())
-> Either ByteBufferException BBRef -> IO ()
forall a b. (a -> b) -> a -> b
$
            ByteBufferException -> Either ByteBufferException BBRef
forall a b. a -> Either a b
Left (String -> String -> ByteBufferException
ByteBufferException String
"free" String
"ByteBuffer has explicitly been freed and is no longer valid.")
    Left ByteBufferException
_ -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return () -- the ByteBuffer is either invalid or has already been freed.

-- | Perform some action with a bytebuffer, with automatic allocation
-- and deallocation.
with :: (MonadIO m, MonadBaseControl IO m)
     => Maybe Int
     -- ^ Initial length of the 'ByteBuffer'.  If 'Nothing', use the
     -- default value of 4MB.
     -> (ByteBuffer -> m a)
     -> m a
with :: forall (m :: * -> *) a.
(MonadIO m, MonadBaseControl IO m) =>
Maybe Int -> (ByteBuffer -> m a) -> m a
with Maybe Int
l ByteBuffer -> m a
action =
  m ByteBuffer -> (ByteBuffer -> m ()) -> (ByteBuffer -> m a) -> m a
forall (m :: * -> *) a b c.
MonadBaseControl IO m =>
m a -> (a -> m b) -> (a -> m c) -> m c
bracket
    (Maybe Int -> m ByteBuffer
forall (m :: * -> *). MonadIO m => Maybe Int -> m ByteBuffer
new Maybe Int
l)
    ByteBuffer -> m ()
forall (m :: * -> *). MonadIO m => ByteBuffer -> m ()
free
    ByteBuffer -> m a
action
{-# INLINE with #-}

-- | Reset a 'BBRef', i.e. copy all the bytes that have not yet
-- been consumed to the front of the buffer.
{-@ resetBBRef :: b:BBRef -> IO {v:BBRef | consumed v == 0 && contained v == contained b - consumed b && size v == size b} @-}
resetBBRef :: BBRef -> IO BBRef
resetBBRef :: BBRef -> IO BBRef
resetBBRef BBRef
bbref = do
    let available :: Int
available = BBRef -> Int
contained BBRef
bbref Int -> Int -> Int
forall a. Num a => a -> a -> a
- BBRef -> Int
consumed BBRef
bbref
    Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
moveBytes (BBRef -> Ptr Word8
ptr BBRef
bbref) (BBRef -> Ptr Word8
ptr BBRef
bbref Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` BBRef -> Int
consumed BBRef
bbref) Int
available
    BBRef -> IO BBRef
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return BBRef { size :: Int
size = BBRef -> Int
size BBRef
bbref
                 , contained :: Int
contained = Int
available
                 , consumed :: Int
consumed = Int
0
                 , ptr :: Ptr Word8
ptr = BBRef -> Ptr Word8
ptr BBRef
bbref
                 }

-- | Make sure the buffer is at least @minSize@ bytes long.
--
-- In order to avoid having to enlarge the buffer too often, we
-- multiply its size by a factor of 1.5 until it is at least @minSize@
-- bytes long.
{-@ enlargeBBRef :: b:BBRef -> i:Nat -> IO {v:BBRef | size v >= i && contained v == contained b && consumed v == consumed b} @-}
enlargeBBRef :: BBRef -> Int -> IO BBRef
enlargeBBRef :: BBRef -> Int -> IO BBRef
enlargeBBRef BBRef
bbref Int
minSize = do
        let getNewSize :: Int -> Int
getNewSize Int
s | Int
s Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
minSize = Int
s
            getNewSize Int
s = Int -> Int
getNewSize (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ (Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling (Double -> Int) -> (Int -> Double) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Double -> Double -> Double
forall a. Num a => a -> a -> a
*(Double
1.5 :: Double)) (Double -> Double) -> (Int -> Double) -> Int -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
s)
            newSize :: Int
newSize = Int -> Int
getNewSize (BBRef -> Int
size BBRef
bbref)
        -- possible optimisation: since reallocation might copy the
        -- bytes anyway, we could discard the consumed bytes,
        -- basically 'reset'ting the buffer on the fly.
        ptr' <- Ptr Word8 -> Int -> IO (Ptr Word8)
forall a. Ptr a -> Int -> IO (Ptr a)
Alloc.reallocBytes (BBRef -> Ptr Word8
ptr BBRef
bbref) Int
newSize
        return BBRef { size = newSize
                     , contained = contained bbref
                     , consumed = consumed bbref
                     , ptr = ptr'
                     }

-- | Copy the contents of a 'ByteString' to a 'ByteBuffer'.
--
-- If necessary, the 'ByteBuffer' is enlarged and/or already consumed
-- bytes are dropped.
copyByteString :: MonadIO m => ByteBuffer -> ByteString -> m ()
copyByteString :: forall (m :: * -> *). MonadIO m => ByteBuffer -> ByteString -> m ()
copyByteString ByteBuffer
bb ByteString
bs =
    String -> ByteBuffer -> (BBRef -> IO ()) -> m ()
forall (m :: * -> *) a.
MonadIO m =>
String -> ByteBuffer -> (BBRef -> IO a) -> m a
bbHandler String
"copyByteString" ByteBuffer
bb BBRef -> IO ()
go
  where
    go :: BBRef -> IO ()
go BBRef
bbref = do
        let (ForeignPtr Word8
bsFptr, Int
bsOffset, Int
bsSize) = ByteString -> (ForeignPtr Word8, Int, Int)
BS.toForeignPtr ByteString
bs
        -- if the byteBuffer is too small, resize it.
        let available :: Int
available = BBRef -> Int
contained BBRef
bbref Int -> Int -> Int
forall a. Num a => a -> a -> a
- BBRef -> Int
consumed BBRef
bbref -- bytes not yet consumed
        bbref' <- if BBRef -> Int
size BBRef
bbref Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
bsSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
available
                  then BBRef -> Int -> IO BBRef
enlargeBBRef BBRef
bbref (Int
bsSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
available)
                  else BBRef -> IO BBRef
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return BBRef
bbref
        -- if it is currently too full, reset it
        bbref'' <- if bsSize + contained bbref' > size bbref'
                   then resetBBRef bbref'
                   else return bbref'
        -- now we can safely copy.
        withForeignPtr bsFptr $ \ Ptr Word8
bsPtr ->
            Ptr (ZonkAny 1) -> Ptr (ZonkAny 1) -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes (BBRef -> Ptr Word8
ptr BBRef
bbref'' Ptr Word8 -> Int -> Ptr (ZonkAny 1)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` BBRef -> Int
contained BBRef
bbref'')
            (Ptr Word8
bsPtr Ptr Word8 -> Int -> Ptr (ZonkAny 1)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
bsOffset)
            Int
bsSize
        writeIORef bb $ Right BBRef {
            size = size bbref''
            , contained = contained bbref'' + bsSize
            , consumed = consumed bbref''
            , ptr = ptr bbref''}

#ifndef mingw32_HOST_OS

-- | Will read at most n bytes from the given 'Fd', in a non-blocking
-- fashion. This function is intended to be used with non-blocking 'Socket's,
-- such the ones created by the @network@ package.
--
-- Returns how many bytes could be read non-blockingly.
fillFromFd :: (MonadIO m, Fail.MonadFail m) => ByteBuffer -> Fd -> Int -> m Int
fillFromFd :: forall (m :: * -> *).
(MonadIO m, MonadFail m) =>
ByteBuffer -> Fd -> Int -> m Int
fillFromFd ByteBuffer
bb Fd
sock Int
maxBytes = if Int
maxBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
    then String -> m Int
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"fillFromFd: negative argument (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
maxBytes String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")")
    else String -> ByteBuffer -> (BBRef -> IO Int) -> m Int
forall (m :: * -> *) a.
MonadIO m =>
String -> ByteBuffer -> (BBRef -> IO a) -> m a
bbHandler String
"fillFromFd" ByteBuffer
bb BBRef -> IO Int
go
  where
    go :: BBRef -> IO Int
go BBRef
bbref = do
        (bbref', readBytes) <- Fd -> BBRef -> Int -> IO (BBRef, Int)
fillBBRefFromFd Fd
sock BBRef
bbref Int
maxBytes
        writeIORef bb $ Right bbref'
        return readBytes

{-
Note: I'd like to use these two definitions:

{-@ measure _available @-}
_available :: BBRef -> Int
_available BBRef{..} = contained - consumed

{-@ measure _free @-}
_free :: BBRef -> Int
_free BBRef{..} = size - contained

but for some reason when I try to do so it does not work.
-}

{-@ fillBBRefFromFd ::
       Fd -> b0: BBRef
    -> maxBytes: Nat -> IO {v: (BBRef, Nat) | maxBytes >= snd v && contained (fst v) - consumed (fst v) == contained b0 - consumed b0 + snd v}
  @-}
fillBBRefFromFd :: Fd -> BBRef -> Int -> IO (BBRef, Int)
fillBBRefFromFd :: Fd -> BBRef -> Int -> IO (BBRef, Int)
fillBBRefFromFd (Fd CInt
sock) BBRef
bbref0 Int
maxBytes = do
  bbref1 <- BBRef -> IO BBRef
prepareSpace BBRef
bbref0
  go 0 bbref1
  where
    -- We enlarge the buffer directly to be able to contain the maximum number
    -- of bytes since the common pattern for this function is to know how many
    -- bytes we need to read -- so we'll eventually fill the enlarged buffer.
    {-@ prepareSpace :: b: BBRef -> IO {v: BBRef | size v - contained v >= maxBytes && contained b - consumed b == contained v - consumed v} @-}
    prepareSpace :: BBRef -> IO BBRef
    prepareSpace :: BBRef -> IO BBRef
prepareSpace BBRef
bbref = do
      let space :: Int
space = BBRef -> Int
size BBRef
bbref Int -> Int -> Int
forall a. Num a => a -> a -> a
- BBRef -> Int
contained BBRef
bbref
      if Int
space Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxBytes
        then if BBRef -> Int
consumed BBRef
bbref Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
          then BBRef -> IO BBRef
prepareSpace (BBRef -> IO BBRef) -> IO BBRef -> IO BBRef
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< BBRef -> IO BBRef
resetBBRef BBRef
bbref
          else BBRef -> Int -> IO BBRef
enlargeBBRef BBRef
bbref (BBRef -> Int
contained BBRef
bbref Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
maxBytes)
        else BBRef -> IO BBRef
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return BBRef
bbref

    {-@ go ::
           readBytes: {v: Nat | v <= maxBytes}
        -> b: {b: BBRef | size b - contained b >= maxBytes - readBytes}
        -> IO {v: (BBRef, Nat) | maxBytes >= snd v && snd v >= readBytes && size (fst v) - contained (fst v) >= maxBytes - snd v && contained (fst v) - consumed (fst v) == (contained b - consumed b) + (snd v - readBytes)}
      @-}
    go :: Int -> BBRef -> IO (BBRef, Int)
    go :: Int -> BBRef -> IO (BBRef, Int)
go Int
readBytes bbref :: BBRef
bbref@BBRef{Int
Ptr Word8
size :: BBRef -> Int
contained :: BBRef -> Int
consumed :: BBRef -> Int
ptr :: BBRef -> Ptr Word8
size :: Int
contained :: Int
consumed :: Int
ptr :: Ptr Word8
..} = if Int
readBytes Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxBytes
      then (BBRef, Int) -> IO (BBRef, Int)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (BBRef
bbref, Int
readBytes)
      else do
        bytes <- CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> CInt -> Ptr CChar -> CSize -> CInt -> IO CInt
c_recv CInt
sock (Ptr (ZonkAny 0) -> Ptr CChar
forall a b. Ptr a -> Ptr b
castPtr (Ptr Word8
ptr Ptr Word8 -> Int -> Ptr (ZonkAny 0)
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
contained)) (Int -> CSize
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
maxBytes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
readBytes)) CInt
0
        if bytes == -1
          then do
            err <- CE.getErrno
            if err == CE.eAGAIN || err == CE.eWOULDBLOCK
              then return (bbref, readBytes)
              else throwIO $ CE.errnoToIOError "ByteBuffer.fillBBRefFromFd: " err Nothing Nothing
          else do
            let bbref' = BBRef
bbref{ contained = contained + bytes }
            go (readBytes + bytes) bbref'

foreign import ccall unsafe "recv"
  -- c_recv returns -1 in the case of errors.
  {-@ assume c_recv :: CInt -> Ptr CChar -> size: {v: CSize | v >= 0} -> flags: CInt -> IO {read: CInt | read >= -1 && size >= read} @-}
  c_recv :: CInt -> Ptr CChar -> CSize -> CInt -> IO CInt

#endif

-- | Try to get a pointer to @n@ bytes from the 'ByteBuffer'.
--
-- Note that the pointer should be used before any other actions are
-- performed on the 'ByteBuffer'. It points to some address within the
-- buffer, so operations such as enlarging the buffer or feeding it
-- new data will change the data the pointer points to.  This is why
-- this function is called unsafe.
{-@ unsafeConsume :: MonadIO m => ByteBuffer -> n:Nat -> m (Either Int ({v:Ptr Word8 | plen v >= n})) @-}
unsafeConsume :: MonadIO m
        => ByteBuffer
        -> Int
        -- ^ n
        -> m (Either Int (Ptr Word8))
        -- ^ Will be @Left missing@ when there are only @n-missing@
        -- bytes left in the 'ByteBuffer'.
unsafeConsume :: forall (m :: * -> *).
MonadIO m =>
ByteBuffer -> Int -> m (Either Int (Ptr Word8))
unsafeConsume ByteBuffer
bb Int
n =
    String
-> ByteBuffer
-> (BBRef -> IO (Either Int (Ptr Word8)))
-> m (Either Int (Ptr Word8))
forall (m :: * -> *) a.
MonadIO m =>
String -> ByteBuffer -> (BBRef -> IO a) -> m a
bbHandler String
"unsafeConsume" ByteBuffer
bb BBRef -> IO (Either Int (Ptr Word8))
forall {b}. BBRef -> IO (Either Int (Ptr b))
go
  where
    go :: BBRef -> IO (Either Int (Ptr b))
go BBRef
bbref = do
        let available :: Int
available = BBRef -> Int
contained BBRef
bbref Int -> Int -> Int
forall a. Num a => a -> a -> a
- BBRef -> Int
consumed BBRef
bbref
        if Int
available Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n
            then Either Int (Ptr b) -> IO (Either Int (Ptr b))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Int (Ptr b) -> IO (Either Int (Ptr b)))
-> Either Int (Ptr b) -> IO (Either Int (Ptr b))
forall a b. (a -> b) -> a -> b
$ Int -> Either Int (Ptr b)
forall a b. a -> Either a b
Left (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
available)
            else do
                ByteBuffer -> Either ByteBufferException BBRef -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef ByteBuffer
bb (Either ByteBufferException BBRef -> IO ())
-> Either ByteBufferException BBRef -> IO ()
forall a b. (a -> b) -> a -> b
$ BBRef -> Either ByteBufferException BBRef
forall a b. b -> Either a b
Right BBRef
bbref { consumed = consumed bbref + n }
                Either Int (Ptr b) -> IO (Either Int (Ptr b))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Int (Ptr b) -> IO (Either Int (Ptr b)))
-> Either Int (Ptr b) -> IO (Either Int (Ptr b))
forall a b. (a -> b) -> a -> b
$ Ptr b -> Either Int (Ptr b)
forall a b. b -> Either a b
Right (BBRef -> Ptr Word8
ptr BBRef
bbref Ptr Word8 -> Int -> Ptr b
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` BBRef -> Int
consumed BBRef
bbref)

-- | As `unsafeConsume`, but instead of returning a `Ptr` into the
-- contents of the `ByteBuffer`, it returns a `ByteString` containing
-- the next @n@ bytes in the buffer.  This involves allocating a new
-- 'ByteString' and copying the @n@ bytes to it.
{-@ consume :: MonadIO m => ByteBuffer -> Nat -> m (Either Int ByteString) @-}
consume :: MonadIO m
        => ByteBuffer
        -> Int
        -> m (Either Int ByteString)
consume :: forall (m :: * -> *).
MonadIO m =>
ByteBuffer -> Int -> m (Either Int ByteString)
consume ByteBuffer
bb Int
n = do
    mPtr <- ByteBuffer -> Int -> m (Either Int (Ptr Word8))
forall (m :: * -> *).
MonadIO m =>
ByteBuffer -> Int -> m (Either Int (Ptr Word8))
unsafeConsume ByteBuffer
bb Int
n
    case mPtr of
        Right Ptr Word8
ptr -> do
            bs <- IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> m ByteString) -> IO ByteString -> m ByteString
forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> Int -> IO ByteString
createBS Ptr Word8
ptr Int
n
            return (Right bs)
        Left Int
missing -> Either Int ByteString -> m (Either Int ByteString)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Either Int ByteString
forall a b. a -> Either a b
Left Int
missing)

{-@ createBS :: p:(Ptr Word8) -> {v:Nat | v <= plen p} -> IO ByteString @-}
createBS :: Ptr Word8 -> Int -> IO ByteString
createBS :: Ptr Word8 -> Int -> IO ByteString
createBS Ptr Word8
ptr Int
n = do
  fp  <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
mallocForeignPtrBytes Int
n
  withForeignPtr fp (\Ptr Word8
p -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
p Ptr Word8
ptr Int
n)
  return (BS.PS fp 0 n)

-- below are liquid haskell qualifiers, and specifications for external functions.

{-@ qualif FPLenPLen(v:Ptr a, fp:ForeignPtr a): fplen fp = plen v @-}

{-@ Foreign.Marshal.Alloc.mallocBytes :: l:Nat -> IO (PtrN a l) @-}
{-@ Foreign.Marshal.Alloc.reallocBytes :: Ptr a -> l:Nat -> IO (PtrN a l) @-}
{-@ assume mallocForeignPtrBytes :: n:Nat -> IO (ForeignPtrN a n) @-}
{-@ type ForeignPtrN a N = {v:ForeignPtr a | fplen v = N} @-}
{-@ Foreign.Marshal.Utils.copyBytes :: p:Ptr a -> q:Ptr a -> {v:Nat | v <= plen p && v <= plen q} -> IO ()@-}
{-@ Foreign.Marshal.Utils.moveBytes :: p:Ptr a -> q:Ptr a -> {v:Nat | v <= plen p && v <= plen q} -> IO ()@-}
{-@ Foreign.Ptr.plusPtr :: p:Ptr a -> n:Nat -> {v:Ptr b | plen v == (plen p) - n} @-}

-- writing down the specification for ByteString is not as straightforward as it seems at first: the constructor
--
-- PS (ForeignPtr Word8) Int Int
--
-- has actually four arguments after unboxing (the ForeignPtr is an
-- Addr# and ForeignPtrContents), so restriciting the length of the
-- ForeignPtr directly in the specification of the datatype does not
-- work.  Instead, I chose to write a specification for toForeignPtr.
-- It seems that the liquidhaskell parser has problems with variables
-- declared in a tuple, so I have to define the following measures to
-- get at the ForeignPtr, offset, and length.
--
-- This is a bit awkward, maybe there is an easier way.

_get1 :: (a,b,c) -> a
_get1 :: forall a b c. (a, b, c) -> a
_get1 (a
x,b
_,c
_) = a
x
{-@ measure _get1 @-}
_get2 :: (a,b,c) -> b
_get2 :: forall a b c. (a, b, c) -> b
_get2 (a
_,b
x,c
_) = b
x
{-@ measure _get2 @-}
_get3 :: (a,b,c) -> c
_get3 :: forall a b c. (a, b, c) -> c
_get3 (a
_,b
_,c
x) = c
x
{-@ measure _get3 @-}

{-@ Data.ByteString.Internal.toForeignPtr :: ByteString ->
  {v:(ForeignPtr Word8, Int, Int) | _get2 v >= 0
                                 && _get2 v <= (fplen (_get1 v))
                                 && _get3 v >= 0
                                 && ((_get3 v) + (_get2 v)) <= (fplen (_get1 v))} @-}