Internal.hs

{-# LANGUAGE PatternSynonyms #-}

-- | Internal module. Not part of the public API.
module Data.PackStream.Get.Internal
  ( getNull, tryNull
  , getBoolean, tryBoolean

  , getFloat, tryFloat

  , getString, tryString
  , getBytes, tryBytes

  , getList, tryList
  , getDictionary, tryDictionary
  , tryStructure
  ) where

import           Compat.Prelude

import qualified Data.ByteString       as S
import qualified Data.Text             as T
import qualified Data.Text.Encoding    as T
import qualified Data.Vector           as V
import qualified Data.HashMap.Lazy     as H
import Data.Hashable (Hashable)

import           Compat.Binary
import           Data.PackStream.Tags

mkGet :: (Word8 -> t -> Get a -> Get b) -> t -> String -> Get b
mkGet tryT f n = do { tag <- getWord8; tryT tag f (fail n) }

-- | Decode a PackStream null value.
getNull :: Get ()
getNull = mkGet tryNull id "expected PackStream null"

-- | Decode a PackStream boolean value.
getBoolean :: Get Bool
getBoolean = mkGet tryBoolean id "expected PackStream boolean"

-- | Decode a PackStream 64-bit float.
getFloat :: Get Double
getFloat = mkGet tryFloat id "expected PackStream float"

-- | Decode a PackStream UTF-8 string.
getString :: Get T.Text
getString = mkGet tryString id "expected PackStream string"

-- | Decode a PackStream byte array.
getBytes :: Get S.ByteString
getBytes = mkGet tryBytes id "expected PackStream bytes"

-- | Decode a PackStream list using the given element decoder.
getList :: Get a -> Get (V.Vector a)
getList g = mkGet (tryList g) id "expected PackStream list"

-- | Decode a PackStream dictionary using the given key and value decoders.
getDictionary :: Hashable a => Get a -> Get b -> Get (H.HashMap a b)
getDictionary k v = mkGet (tryDictionary k v) id "Dictionary"


----------------------------------------------------------------------------
-- primitives that take a tag as first argument

-- | Try to decode null given a tag byte; apply @f@ on success or fall through to the continuation.
{-# INLINE tryNull #-}
tryNull :: Word8 -> (() -> a) -> Get a -> Get a
tryNull tag f cont = case tag of
  TAG_Null -> pure $! f ()
  _       -> cont

-- | Try to decode a boolean given a tag byte; apply @f@ on success or fall through to the continuation.
{-# INLINE tryBoolean #-}
tryBoolean :: Word8 -> (Bool -> a) -> Get a -> Get a
tryBoolean tag f cont = case tag of
  TAG_false -> pure $! f False
  TAG_true  -> pure $! f True
  _         -> cont


-- | Try to decode a float given a tag byte; apply @f@ on success or fall through to the continuation.
{-# INLINE tryFloat #-}
tryFloat :: Word8 -> (Double -> a) -> Get a -> Get a
tryFloat tag f cont = case tag of
  TAG_Float -> f <$> getFloat64be
  _           -> cont

-- | Try to decode a string given a tag byte; apply @f@ on success or fall through to the continuation.
{-# INLINE tryString #-}
tryString :: Word8 -> (T.Text -> a) -> Get a -> Get a
tryString tag f cont = case tag of
    t | Just sz <- is_TAG_STRING_SHORT t -> cont' sz
    TAG_STRING_8                       -> cont' . intCast =<< getWord8
    TAG_STRING_16                      -> cont' . intCast =<< getWord16be
    TAG_STRING_32                      -> cont' =<< getWord32be
    _                              -> cont
  where
    cont' len = do
      len' <- fromSizeM "tryString: data exceeds capacity of ByteString/Text" len
      bs <- getByteString len'
      case T.decodeUtf8' bs of
        Left _  -> fail "tryString: invalid UTF-8 encoding"
        Right v -> pure $! f v

-- | Try to decode a byte array given a tag byte; apply @f@ on success or fall through to the continuation.
{-# INLINE tryBytes #-}
tryBytes :: Word8 -> (S.ByteString -> a) -> Get a -> Get a
tryBytes tag f cont = case tag of
    TAG_BYTE_ARRAY_8  -> cont' . intCast =<< getWord8
    TAG_BYTE_ARRAY_16 -> cont' . intCast =<< getWord16be
    TAG_BYTE_ARRAY_32 -> cont' =<< getWord32be
    _         -> cont
  where
    cont' len = do
      len' <- fromSizeM "tryBytes: data exceeds capacity of ByteString" len
      f <$> getByteString len'

-- | Try to decode a list given a tag byte; apply @f@ on success or fall through to the continuation.
{-# INLINE tryList #-}
tryList :: Get b -> Word8 -> (V.Vector b -> a) -> Get a -> Get a
tryList g tag f cont = case tag of
    t | Just sz <- is_TAG_LIST_SHORT t -> cont' sz
    TAG_LIST_8                       -> cont' . intCast =<< getWord8
    TAG_LIST_16                      -> cont' . intCast =<< getWord16be
    TAG_LIST_32                      -> cont' =<< getWord32be
    _                                -> cont
  where
    cont' len = do
      len' <- fromSizeM "tryList: data exceeds capacity of Vector" len
      f <$> V.replicateM len' g

-- | Try to decode a dictionary given a tag byte; apply @f@ on success or fall through to the continuation.
{-# INLINE tryDictionary #-}
tryDictionary :: Hashable k => Get k -> Get v -> Word8 -> (H.HashMap k v -> a) -> Get a -> Get a
tryDictionary k v tag f cont = case tag of
    t | Just sz <- is_TAG_DICTIONARY_SHORT t -> cont' sz
    TAG_DICTIONARY_8               -> cont' . intCast =<< getWord8
    TAG_DICTIONARY_16              -> cont' . intCast =<< getWord16be
    TAG_DICTIONARY_32              -> cont' =<< getWord32be
    _                              -> cont
  where
    cont' len = do
      len' <- fromSizeM "tryDictionary: data exceeds capacity of Vector" len
      fmap (f . H.fromList) $ replicateM len' ((,) <$> k <*> v)
      -- fmap f $ V.replicateM len' ((,) <$> k <*> v)

-- {-# INLINE tryStructure #-}
-- tryStructure :: Word8 -> (H.HashMap T.Text Ps -> a) -> Get a -> Get a
-- tryStructure tag f cont = case tag of
--   TAG_float64 -> f <$> getFloat64be
--   _           -> cont

-- | Try to decode a structure given a tag byte; returns (structTag, fields) via @f@ on success.
{-# INLINE tryStructure #-}
tryStructure :: Get b -> Word8 -> ((Word8, V.Vector b) -> a) -> Get a -> Get a
tryStructure g tag f cont = case is_TAG_STRUCTURE tag of
  Just nfields -> do
    structTag <- getWord8
    nfields' <- fromSizeM "tryStructure: data exceeds capacity of Vector" nfields
    fields <- V.replicateM nfields' g
    pure $! f (structTag, fields)
  Nothing -> cont

fromSizeM :: String -> Word32 -> Get Int
fromSizeM label sz = maybe (fail label) pure (intCastMaybe sz)