-- | Accumulate strict bytestrings
--
-- Intended for qualified import.
--
-- > import Network.GRPC.Util.AccumulatedByteString (AccumulatedByteString)
-- > import Network.GRPC.Util.AccumulatedByteString qualified as AccBS
module Network.GRPC.Util.AccumulatedByteString (
    AccumulatedByteString -- opaque
    -- * Construction
  , init
  , append
    -- * Query
  , length
    -- * Destruction
  , splitAt
  ) where

import Prelude hiding (init, splitAt, length)

import Data.ByteString qualified as BS.Strict
import Data.ByteString qualified as Strict (ByteString)

{-------------------------------------------------------------------------------
  Definition
-------------------------------------------------------------------------------}

-- | Accumulate strict bytestrings until we have enough bytes
data AccumulatedByteString = AccBS {
       -- | Total accumulated length
       AccumulatedByteString -> Word
accLength  :: !Word

       -- | Accumulated bytestrings, in reverse order
     , AccumulatedByteString -> [ByteString]
accStrings :: [Strict.ByteString]
     }

{-------------------------------------------------------------------------------
  Construction
-------------------------------------------------------------------------------}

init :: Maybe Strict.ByteString -> AccumulatedByteString
init :: Maybe ByteString -> AccumulatedByteString
init Maybe ByteString
Nothing   = Word -> [ByteString] -> AccumulatedByteString
AccBS Word
0 []
init (Just ByteString
bs) = Word -> [ByteString] -> AccumulatedByteString
AccBS (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Word) -> Int -> Word
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.Strict.length ByteString
bs) [ByteString
bs]

append :: AccumulatedByteString -> Strict.ByteString -> AccumulatedByteString
append :: AccumulatedByteString -> ByteString -> AccumulatedByteString
append AccBS{Word
accLength :: AccumulatedByteString -> Word
accLength :: Word
accLength, [ByteString]
accStrings :: AccumulatedByteString -> [ByteString]
accStrings :: [ByteString]
accStrings} ByteString
bs = AccBS{
      accLength :: Word
accLength  = Word
accLength Word -> Word -> Word
forall a. Num a => a -> a -> a
+ Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
BS.Strict.length ByteString
bs)
    , accStrings :: [ByteString]
accStrings = ByteString
bs ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
accStrings
    }

{-------------------------------------------------------------------------------
  Query
-------------------------------------------------------------------------------}

length :: AccumulatedByteString -> Word
length :: AccumulatedByteString -> Word
length = AccumulatedByteString -> Word
accLength

{-------------------------------------------------------------------------------
  Destruction
-------------------------------------------------------------------------------}

-- | Split the accumulated bytestrings after @n@ bytes
--
-- Preconditions:
--
-- 1. The accumulated length must be @>= n@
-- 2. The accumulated length without the most recently added bytestring
--    must be @<= n@
--
-- These two preconditions together imply that all accumulated bytestrings
-- except the most recent will be required.
splitAt ::
     Word -- ^ @n@
  -> AccumulatedByteString
  -> (Strict.ByteString, Maybe Strict.ByteString)
splitAt :: Word -> AccumulatedByteString -> (ByteString, Maybe ByteString)
splitAt Word
n AccBS{Word
accLength :: AccumulatedByteString -> Word
accLength :: Word
accLength, [ByteString]
accStrings :: AccumulatedByteString -> [ByteString]
accStrings :: [ByteString]
accStrings} =
    if Word
leftoverLen Word -> Word -> Bool
forall a. Eq a => a -> a -> Bool
== Word
0 then
      ([ByteString] -> ByteString
BS.Strict.concat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse [ByteString]
accStrings, Maybe ByteString
forall a. Maybe a
Nothing)
    else
      -- If @leftoverLen > 0@ then @accLength > 0@: we must have some strings
      case [ByteString]
accStrings of
        []              -> [Char] -> (ByteString, Maybe ByteString)
forall a. HasCallStack => [Char] -> a
error [Char]
"splitAt: invalid AccumulatedByteString"
        ByteString
mostRecent:[ByteString]
rest ->
          let neededLen :: Int
neededLen =
                -- cannot underflow due to precondition (2)
                ByteString -> Int
BS.Strict.length ByteString
mostRecent Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
leftoverLen
              (ByteString
needed, ByteString
leftover) =
                Int -> ByteString -> (ByteString, ByteString)
BS.Strict.splitAt Int
neededLen ByteString
mostRecent
          in ( [ByteString] -> ByteString
BS.Strict.concat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse (ByteString
needed ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
rest)
             , if ByteString -> Bool
BS.Strict.null ByteString
leftover
                 then Maybe ByteString
forall a. Maybe a
Nothing
                 else ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
leftover
             )
  where
    -- cannot underflow due to precondition (1)
    leftoverLen :: Word
    leftoverLen :: Word
leftoverLen = Word
accLength Word -> Word -> Word
forall a. Num a => a -> a -> a
- Word
n