Ps.hs

{-# LANGUAGE IncoherentInstances  #-}
{-# LANGUAGE OverloadedLists      #-}

-- | Core PackStream AST: the 'Ps' type and the 'PackStream' type class for
-- converting between Haskell types and PackStream values.
--
-- Specification: <https://neo4j.com/docs/bolt/current/packstream/>
module Data.PackStream.Ps
    ( -- * Core types
      Tag
    , Ps(..)
    , structFromDict
    , structureSingleton
      -- * PackStream class
    , PackStream(..)
      -- * Operators
    , (.:), (.=)
      -- * Accessors
    , getPs, putPs
      -- * Type matching
    , typeMismatch
    , withNull, withBoolean, withInteger, withFloat
    , withBytes, withString, withList, withDictionary
    ) where

import           Compat.Prelude
import           Prelude                       hiding (putStr)

import qualified Data.ByteString               as S
import qualified Data.ByteString.Lazy          as L
import qualified Data.ByteString.Short         as SBS
import qualified Data.HashMap.Strict           as HashMap
import           Data.Kind                     (Constraint, Type)
import           Data.List.NonEmpty            (NonEmpty)
import qualified Data.List.NonEmpty            as NEL
import qualified Data.Text                     as T
import qualified Data.Text.Lazy                as LT
import           Data.Typeable
import qualified Data.Vector                   as V
import qualified Data.HashMap.Lazy             as H

import           Data.PackStream.Get.Internal
import           Data.PackStream.Integer
import           Data.PackStream.Put
import           Data.PackStream.Result
import           Data.PackStream.Tags ()

import           Compat.Binary
import           TextShow (showt)


-- | Structure tag byte identifying a PackStream structure type.
type Tag :: Type
type Tag = Word8

-- | A PackStream value. This is the intermediate AST used for serialization.
type Ps :: Type
data Ps
  = PsNull
    -- ^ missing or empty value
  | PsBoolean !Bool
    -- ^ true or false
  | PsInteger !PSInteger
    -- ^ signed 64-bit integer
  | PsFloat !Double
    -- ^ 64-bit floating point number
  | PsBytes !S.ByteString
    -- ^ byte array
  | PsString !T.Text
    -- ^ unicode text, UTF-8
  | PsList !(V.Vector Ps)
    -- ^ ordered collection of values
  | PsDictionary !(H.HashMap T.Text Ps)
    -- ^ collection of key-value entries (no order guaranteed)
  | PsStructure !Tag !(V.Vector Ps)
    -- ^ composite value with a type signature
    --
    -- fields being a Vector is a bit wasteful, but the spec demands it
    -- in practice there is always just 1 field, which is a dictionary
    -- Control messages all use dictionaries: https://neo4j.com/docs/bolt/current/bolt/message/#messages
    -- Datatypes all use dictionaries: https://neo4j.com/docs/bolt/current/bolt/structure-semantics/
  deriving stock (Show, Eq)

-- | Build a structure with a single dictionary field.
structFromDict :: Tag -> H.HashMap T.Text Ps -> Ps
structFromDict tag map = PsStructure tag $ V.singleton $ PsDictionary map

-- | Build a structure with exactly one field.
structureSingleton :: Tag -> Ps -> Ps
structureSingleton tag ps = PsStructure tag $ V.singleton ps


-- | Look up a key in a 'PsDictionary' and decode the value.
(.:) :: PackStream a => Ps -> T.Text -> Result a
(PsDictionary m) .: key = case H.lookup key m of
  Nothing -> Error $ "missing key \"" <> key <> "\""
  Just v -> fromPs v
m .: _ = Error $ ("expected PsDictionary got " :: T.Text) <> (showt . typeOf $ m)

-- | Build a key-value pair for constructing a 'PsDictionary'.
(.=) :: PackStream a => T.Text -> a -> (Ps, Ps)
k .= a = (PsString k, toPs a)

instance NFData Ps where
  rnf obj = case obj of
    PsList a -> rnf a
    PsDictionary   m -> rnf m
    _             -> ()

-- | Decode any PackStream value from the binary stream.
getPs :: Get Ps
getPs = do
  tag <- getWord8

  tryNull tag (const PsNull) $
    tryBoolean    tag PsBoolean $
    tryPSInteger  tag PsInteger $
    tryFloat      tag PsFloat $
    tryString     tag PsString $
    tryBytes      tag PsBytes $
    tryList getPs tag PsList $
    tryDictionary getString getPs tag PsDictionary $
    tryStructure getPs tag (uncurry PsStructure) $
    fail ("getPs: internal error " ++ show tag) -- should never happen

-- | Encode any PackStream value to the binary stream.
putPs :: Ps -> Put
putPs = \case
  PsNull         -> putNull
  PsBoolean    b -> putBoolean b
  PsInteger    n -> put n
  PsFloat      d -> putFloat d
  PsString     t -> putString t
  PsBytes      b -> putBytes b
  PsList       a -> putList putPs a
  PsDictionary m -> putDictionary putString putPs m
  PsStructure t fs -> do
    let nfields = V.length fs
    putWord8 (0xB0 .|. fromIntegral nfields)
    putWord8 t
    V.mapM_ putPs fs

-- | This 'Persist' instance encodes\/decodes to\/from PackStream format
instance Persist Ps where
  get = getPs
  put = putPs

-- | Class for converting between PackStream 'Ps's and native Haskell types.
type PackStream :: Type -> Constraint
class PackStream a where
  toPs   :: a -> Ps

  -- | Encodes directly to 'Put' monad bypassing the intermediate 'Ps' AST
  --
  -- @since 1.1.0.0
  toBinary :: a -> Put
  toBinary = putPs . toPs

  fromPs :: Ps -> Result a

-- core instances

-- | The trivial identity 'PackStream' instance
instance PackStream Ps where
  toPs = id
  toBinary = putPs
  fromPs = pure

-- | Encodes as 'PsNull'
instance PackStream () where
  toPs _ = PsNull
  toBinary _ = putNull
  fromPs = withNull "()" (pure ())

instance PackStream Bool where
  toPs = PsBoolean
  toBinary = putBoolean
  fromPs = withBoolean "Boolean" pure

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

instance PackStream PSInteger where
  toPs = PsInteger
  toBinary = put
  fromPs = withInteger "PSInteger" pure

fromPsInteger :: FromPSInteger i => T.Text -> Ps -> Result i
fromPsInteger expected = withInteger expected go
  where
    go j = case fromPSInteger j of
      Just j' -> pure j'
      Nothing -> fail $ T.unpack $ "PackStream integer " <> showt j <> " cannot be decoded into " <> expected

instance PackStream Word32 where
  toPs = PsInteger . toPSInteger
  toBinary = put . toPSInteger
  fromPs = fromPsInteger "Word32"

instance PackStream Word16 where
  toPs = PsInteger . toPSInteger
  toBinary = put . toPSInteger
  fromPs = fromPsInteger "Word16"

instance PackStream Word8 where
  toPs = PsInteger . toPSInteger
  toBinary = put . toPSInteger
  fromPs = fromPsInteger "Word8"

instance PackStream Int64 where
  toPs = PsInteger . toPSInteger
  toBinary = put . toPSInteger
  fromPs = fromPsInteger "Int64"

instance PackStream Int32 where
  toPs = PsInteger . toPSInteger
  toBinary = put . toPSInteger
  fromPs = fromPsInteger "Int32"

instance PackStream Int16 where
  toPs = PsInteger . toPSInteger
  toBinary = put . toPSInteger
  fromPs = fromPsInteger "Int16"

instance PackStream Int8 where
  toPs = PsInteger . toPSInteger
  toBinary = put . toPSInteger
  fromPs = fromPsInteger "Int8"

instance PackStream Int where
  toPs = PsInteger . toPSInteger
  toBinary = put . toPSInteger
  fromPs = fromPsInteger "Int"

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

-- | This instance decodes 64bit and 32bit floats from PackStream streams into a 'Double'
instance PackStream Double where
  toPs = PsFloat
  toBinary = putFloat
  fromPs = withFloat "Double" pure

instance PackStream S.ByteString where
  toPs = PsBytes
  toBinary = putBytes
  fromPs = withBytes "ByteString" pure

-- Because of overlapping instance, this must be above [a]
instance PackStream String where
  toPs = toPs . T.pack
  toBinary = putString . T.pack
  fromPs obj = T.unpack <$> fromPs obj

instance PackStream a => PackStream (V.Vector a) where
  toPs = PsList . V.map toPs
  toBinary = putList toBinary
  fromPs = withList "Vector" (traverse fromPs)



-- | 'Maybe's are encoded as nullable types, i.e. 'Nothing' is encoded as @nil@.
--
-- __NOTE__: Encoding nested 'Maybe's or 'Maybe's enclosing types which encode to @nil@ (such as '()') will break round-tripping
instance PackStream a => PackStream (Maybe a) where
  toPs = \case
    Just a  -> toPs a
    Nothing -> PsNull
  toBinary = \case
    Just a  -> toBinary a
    Nothing -> putNull

  fromPs = \case
    PsNull -> pure Nothing
    obj       -> Just <$> fromPs obj

-- UTF8 string like

instance PackStream L.ByteString where
  toPs = PsBytes . L.toStrict
  toBinary = putBytes . L.toStrict
  fromPs obj = L.fromStrict <$> fromPs obj

instance PackStream SBS.ShortByteString where
  toPs = PsBytes . SBS.fromShort
  toBinary = putBytes . SBS.fromShort
  fromPs obj = SBS.toShort <$> fromPs obj

instance PackStream T.Text where
  toPs = PsString
  toBinary = putString
  fromPs = withString "Text" pure

instance PackStream LT.Text where
  toPs = toPs . LT.toStrict
  toBinary = putString . LT.toStrict
  fromPs obj = LT.fromStrict <$> fromPs obj

-- array like

instance PackStream a => PackStream [a] where
  toPs = toPs . V.fromList
  toBinary = putList toBinary . V.fromList
  fromPs obj = V.toList <$> fromPs obj

instance PackStream a => PackStream (NonEmpty a) where
  toPs = toPs . NEL.toList
  toBinary = toBinary . NEL.toList
  fromPs o = do
    lst <- fromPs o
    case NEL.nonEmpty lst of
      Just as -> Success as
      Nothing -> Error "empty list"

-- dictionary like

instance (PackStream v) => PackStream (HashMap.HashMap T.Text v) where
  toPs = PsDictionary . H.map toPs
  toBinary = putDictionary putString toBinary . H.fromList . HashMap.toList
  fromPs = withDictionary "HashMap" (traverse fromPs)



-- tuples

instance (PackStream a1, PackStream a2) => PackStream (a1, a2) where
  toPs (a1, a2) = PsList [toPs a1, toPs a2]
  toBinary (a1, a2) = putList' 2 $ do { toBinary a1; toBinary a2 }
  fromPs (PsList [a1, a2]) = (,) <$> fromPs a1 <*> fromPs a2
  fromPs obj                    = typeMismatch "2-tuple" obj

instance (PackStream a1, PackStream a2, PackStream a3) => PackStream (a1, a2, a3) where
  toPs (a1, a2, a3) = PsList [toPs a1, toPs a2, toPs a3]
  toBinary (a1, a2, a3) = putList' 3 $ do { toBinary a1; toBinary a2; toBinary a3 }
  fromPs (PsList [a1, a2, a3]) = (,,) <$> fromPs a1 <*> fromPs a2 <*> fromPs a3
  fromPs obj = typeMismatch "3-tuple" obj

instance (PackStream a1, PackStream a2, PackStream a3, PackStream a4) => PackStream (a1, a2, a3, a4) where
  toPs (a1, a2, a3, a4) = PsList [toPs a1, toPs a2, toPs a3, toPs a4]
  toBinary (a1, a2, a3, a4) = putList' 4 $ do { toBinary a1; toBinary a2; toBinary a3; toBinary a4 }
  fromPs (PsList [a1, a2, a3, a4]) = (,,,) <$> fromPs a1 <*> fromPs a2 <*> fromPs a3 <*> fromPs a4
  fromPs obj = typeMismatch "4-tuple" obj

instance (PackStream a1, PackStream a2, PackStream a3, PackStream a4, PackStream a5) => PackStream (a1, a2, a3, a4, a5) where
  toPs (a1, a2, a3, a4, a5) = PsList [toPs a1, toPs a2, toPs a3, toPs a4, toPs a5]
  toBinary (a1, a2, a3, a4, a5) = putList' 5 $ do { toBinary a1; toBinary a2; toBinary a3; toBinary a4; toBinary a5 }
  fromPs (PsList [a1, a2, a3, a4, a5]) = (,,,,) <$> fromPs a1 <*> fromPs a2 <*> fromPs a3 <*> fromPs a4 <*> fromPs a5
  fromPs obj = typeMismatch "5-tuple" obj

instance (PackStream a1, PackStream a2, PackStream a3, PackStream a4, PackStream a5, PackStream a6) => PackStream (a1, a2, a3, a4, a5, a6) where
  toPs (a1, a2, a3, a4, a5, a6) = PsList [toPs a1, toPs a2, toPs a3, toPs a4, toPs a5, toPs a6]
  toBinary (a1, a2, a3, a4, a5, a6) = putList' 6 $ do { toBinary a1; toBinary a2; toBinary a3; toBinary a4; toBinary a5; toBinary a6 }
  fromPs (PsList [a1, a2, a3, a4, a5, a6]) = (,,,,,) <$> fromPs a1 <*> fromPs a2 <*> fromPs a3 <*> fromPs a4 <*> fromPs a5 <*> fromPs a6
  fromPs obj = typeMismatch "6-tuple" obj

instance (PackStream a1, PackStream a2, PackStream a3, PackStream a4, PackStream a5, PackStream a6, PackStream a7) => PackStream (a1, a2, a3, a4, a5, a6, a7) where
  toPs (a1, a2, a3, a4, a5, a6, a7) = PsList [toPs a1, toPs a2, toPs a3, toPs a4, toPs a5, toPs a6, toPs a7]
  toBinary (a1, a2, a3, a4, a5, a6, a7) = putList' 7 $ do { toBinary a1; toBinary a2; toBinary a3; toBinary a4; toBinary a5; toBinary a6; toBinary a7 }
  fromPs (PsList [a1, a2, a3, a4, a5, a6, a7]) = (,,,,,,) <$> fromPs a1 <*> fromPs a2 <*> fromPs a3 <*> fromPs a4 <*> fromPs a5 <*> fromPs a6 <*> fromPs a7
  fromPs obj = typeMismatch "7-tuple" obj

instance (PackStream a1, PackStream a2, PackStream a3, PackStream a4, PackStream a5, PackStream a6, PackStream a7, PackStream a8) => PackStream (a1, a2, a3, a4, a5, a6, a7, a8) where
  toPs (a1, a2, a3, a4, a5, a6, a7, a8) = PsList [toPs a1, toPs a2, toPs a3, toPs a4, toPs a5, toPs a6, toPs a7, toPs a8]
  toBinary (a1, a2, a3, a4, a5, a6, a7, a8) = putList' 8 $ do { toBinary a1; toBinary a2; toBinary a3; toBinary a4; toBinary a5; toBinary a6; toBinary a7; toBinary a8 }
  fromPs (PsList [a1, a2, a3, a4, a5, a6, a7, a8]) = (,,,,,,,) <$> fromPs a1 <*> fromPs a2 <*> fromPs a3 <*> fromPs a4 <*> fromPs a5 <*> fromPs a6 <*> fromPs a7 <*> fromPs a8
  fromPs obj = typeMismatch "8-tuple" obj

instance (PackStream a1, PackStream a2, PackStream a3, PackStream a4, PackStream a5, PackStream a6, PackStream a7, PackStream a8, PackStream a9) => PackStream (a1, a2, a3, a4, a5, a6, a7, a8, a9) where
  toPs (a1, a2, a3, a4, a5, a6, a7, a8, a9) = PsList [toPs a1, toPs a2, toPs a3, toPs a4, toPs a5, toPs a6, toPs a7, toPs a8, toPs a9]
  toBinary (a1, a2, a3, a4, a5, a6, a7, a8, a9) = putList' 8 $ do { toBinary a1; toBinary a2; toBinary a3; toBinary a4; toBinary a5; toBinary a6; toBinary a7; toBinary a8; toBinary a9 }
  fromPs (PsList [a1, a2, a3, a4, a5, a6, a7, a8, a9]) = (,,,,,,,,) <$> fromPs a1 <*> fromPs a2 <*> fromPs a3 <*> fromPs a4 <*> fromPs a5 <*> fromPs a6 <*> fromPs a7 <*> fromPs a8 <*> fromPs a9
  fromPs obj = typeMismatch "9-tuple" obj

-- | Report a type mismatch error for the expected type and actual 'Ps' value.
typeMismatch :: T.Text -> Ps -> Result a
typeMismatch expected obj = fail $ T.unpack ("PackStream " <> got <> " type cannot be decoded into " <> expected)
  where
    got = case obj of
      PsNull           -> "nil"
      PsList v         -> "array[" <> showt (V.length v) <> "]"
      PsDictionary v   -> "map[" <> showt (H.size v) <> "]"
      PsString _       -> "str"
      PsBoolean _      -> "bool"
      PsInteger _      -> "int"
      PsFloat _        -> "float"
      PsBytes _        -> "bin"
      PsStructure ty _ -> "structure[" <> showt ty <> "]"

-- | Match a 'PsNull', or report a type mismatch.
withNull :: T.Text -> Result a -> Ps -> Result a
withNull _ f PsNull  = f
withNull expected _ got = typeMismatch expected got

-- | Match a 'PsBoolean', or report a type mismatch.
withBoolean :: T.Text -> (Bool -> Result a) -> Ps -> Result a
withBoolean _ f (PsBoolean b) = f b
withBoolean expected _ got = typeMismatch expected got

-- | Match a 'PsInteger', or report a type mismatch.
withInteger :: T.Text -> (PSInteger -> Result a) -> Ps -> Result a
withInteger _ f (PsInteger i) = f i
withInteger expected _ got = typeMismatch expected got

-- | Match a 'PsFloat', or report a type mismatch.
withFloat :: T.Text -> (Double -> Result a) -> Ps -> Result a
withFloat _ f (PsFloat x) = f x
withFloat expected _ got   = typeMismatch expected got

-- | Match a 'PsBytes', or report a type mismatch.
withBytes :: T.Text -> (S.ByteString -> Result a) -> Ps -> Result a
withBytes _ f (PsBytes i) = f i
withBytes expected _ got = typeMismatch expected got

-- | Match a 'PsString', or report a type mismatch.
withString :: T.Text -> (T.Text -> Result a) -> Ps -> Result a
withString _ f (PsString i) = f i
withString expected _ got = typeMismatch expected got

-- | Match a 'PsList', or report a type mismatch.
withList :: T.Text -> (V.Vector Ps -> Result a) -> Ps -> Result a
withList _ f (PsList xs) = f xs
withList expected _ got   = typeMismatch expected got

-- | Match a 'PsDictionary', or report a type mismatch.
withDictionary :: T.Text -> (H.HashMap T.Text Ps -> Result a) -> Ps -> Result a
withDictionary _ f (PsDictionary xs) = f xs
withDictionary expected _ got = typeMismatch expected got