Serialize.hs

{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- | Binary serializer elements
module RON.Binary.Serialize (
    serialize,
    serializeAtom,
    serializeString,
) where

import RON.Prelude

import Data.Binary qualified as Binary
import Data.Binary.Put (putDoublebe, runPut)
import Data.Bits (bit, shiftL, (.|.))
import Data.ByteString.Lazy (cons, fromStrict)
import Data.ByteString.Lazy qualified as BSL
import Data.Text.Encoding (encodeUtf8)
import Data.ZigZag (zzEncode)

import RON.Binary.Types (Desc (..), Size, descIsOp)
import RON.Types (
    Atom (AFloat, AInteger, AString, AUuid),
    ClosedOp (..),
    Op (..),
    UUID (UUID),
    WireChunk (Closed, Query, Value),
    WireFrame,
    WireReducedChunk (..),
 )
import RON.Util.Word (Word4, b0000, leastSignificant4, safeCast)

-- | Serialize a frame
serialize :: WireFrame -> Either String ByteStringL
serialize chunks = ("RON2" <>) <$> serializeBody
  where
    serializeBody = foldChunks =<< traverse serializeChunk chunks

    chunkSize :: Bool -> Int64 -> Either String ByteStringL
    chunkSize continue x
        | x < bit 31 = Right $ Binary.encode s'
        | otherwise = Left $ "chunk size is too big: " ++ show x
      where
        s = fromIntegral x :: Size
        s'
            | continue = s .|. bit 31
            | otherwise = s

    foldChunks :: [ByteStringL] -> Either String ByteStringL
    foldChunks = \case
        [] -> chunkSize False 0
        [c] -> (<> c) <$> chunkSize False (BSL.length c)
        c : cs ->
            fold
                <$> sequence
                    [chunkSize True (BSL.length c), pure c, foldChunks cs]

-- | Serialize a chunk
serializeChunk :: WireChunk -> Either String ByteStringL
serializeChunk = \case
    Closed op -> serializeClosedOp DOpClosed op
    Value rchunk -> serializeReducedChunk False rchunk
    Query rchunk -> serializeReducedChunk True rchunk

-- | Serialize a closed op
serializeClosedOp :: Desc -> ClosedOp -> Either String ByteStringL
serializeClosedOp desc ClosedOp{..} = do
    keys <-
        sequenceA
            [ serializeUuidReducer reducerId
            , serializeUuidObject objectId
            , serializeUuidOpId opId
            , serializeUuidRef refId
            ]
    payloadValue <- traverse serializeAtom payload
    serializeWithDesc desc $ fold $ keys ++ payloadValue
  where
    Op{opId, refId, payload} = op
    serializeUuidReducer = serializeWithDesc DUuidReducer . serializeUuid
    serializeUuidObject = serializeWithDesc DUuidObject . serializeUuid
    serializeUuidOpId = serializeWithDesc DUuidOp . serializeUuid
    serializeUuidRef = serializeWithDesc DUuidRef . serializeUuid

-- | Serialize a reduced op
serializeReducedOp :: Desc -> UUID -> UUID -> Op -> Either String ByteStringL
serializeReducedOp d reducerId objectId op =
    serializeClosedOp d ClosedOp{reducerId, objectId, op}

-- | Serialize a 'UUID'
serializeUuid :: UUID -> ByteStringL
serializeUuid (UUID x y) = Binary.encode x <> Binary.encode y

-- | Encode descriptor
encodeDesc :: Desc -> Word4
encodeDesc = leastSignificant4 . fromEnum

-- | Prepend serialized bytes with descriptor
serializeWithDesc ::
    Desc ->
    -- | body
    ByteStringL ->
    Either String ByteStringL
serializeWithDesc d body = do
    (lengthDesc, lengthExtended) <- lengthFields
    let descByte = safeCast (encodeDesc d) `shiftL` 4 .|. safeCast lengthDesc
    pure $ descByte `cons` lengthExtended <> body
  where
    len = BSL.length body
    lengthFields = case d of
        DAtomString
            | len == 0 -> Right (b0000, mkLengthExtended)
            | len < 16 -> Right (leastSignificant4 len, BSL.empty)
            | len < bit 31 -> Right (b0000, mkLengthExtended)
            | otherwise -> Left "String is too long"
        _
            | descIsOp d -> Right (b0000, BSL.empty)
            | len < 16 -> Right (leastSignificant4 len, BSL.empty)
            | len == 16 -> Right (b0000, BSL.empty)
            | otherwise -> Left "impossible"
    mkLengthExtended
        | len < 128 = Binary.encode (fromIntegral len :: Word8)
        | otherwise = Binary.encode (fromIntegral len .|. bit 31 :: Word32)

-- | Serialize an 'Atom'
serializeAtom :: Atom -> Either String ByteStringL
serializeAtom = \case
    AFloat f -> serializeWithDesc DAtomFloat $ serializeFloat f
    AInteger i -> serializeWithDesc DAtomInteger $ Binary.encode $ zzEncode64 i
    AString s -> serializeWithDesc DAtomString $ serializeString s
    AUuid u -> serializeWithDesc DAtomUuid $ serializeUuid u
  where
    {-# INLINE zzEncode64 #-}
    zzEncode64 :: Int64 -> Word64
    zzEncode64 = zzEncode

-- | Serialize a float atom
serializeFloat :: Double -> ByteStringL
serializeFloat = runPut . putDoublebe

-- | Serialize a reduced chunk
serializeReducedChunk :: Bool -> WireReducedChunk -> Either String ByteStringL
serializeReducedChunk isQuery WireReducedChunk{..} = do
    header <-
        serializeClosedOp
            (if isQuery then DOpQueryHeader else DOpHeader)
            wrcHeader
    body <- foldMapA (serializeReducedOp DOpReduced reducerId objectId) wrcBody
    pure $ header <> body
  where
    ClosedOp{..} = wrcHeader

-- | Serialize a string atom
serializeString :: Text -> ByteStringL
serializeString = fromStrict . encodeUtf8

foldMapA :: (Applicative f, Foldable t, Monoid b) => (a -> f b) -> t a -> f b
foldMapA f = fmap fold . traverse f . toList