{-# LANGUAGE DataKinds           #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE MonoLocalBinds      #-}
{-# LANGUAGE PolyKinds           #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Network.TypedProtocol.Stateful.Codec.CBOR
  ( module Network.TypedProtocol.Stateful.Codec
  , DeserialiseFailure
  , mkCodecCborLazyBS
  , mkCodecCborStrictBS
  ) where

import Control.Monad.Class.MonadST (MonadST (..))

import Codec.CBOR.Decoding qualified as CBOR (Decoder)
import Codec.CBOR.Encoding qualified as CBOR (Encoding)
import Codec.CBOR.Read qualified as CBOR
import Codec.CBOR.Write qualified as CBOR
import Data.ByteString qualified as BS
import Data.ByteString.Builder qualified as BS
import Data.ByteString.Builder.Extra qualified as BS
import Data.ByteString.Lazy qualified as LBS
import Data.ByteString.Lazy.Internal qualified as LBS (smallChunkSize)

import Network.TypedProtocol.Stateful.Codec
import Network.TypedProtocol.Codec.CBOR (DeserialiseFailure,
         convertCborDecoderBS, convertCborDecoderLBS)
import Network.TypedProtocol.Core


-- | Construct a 'Codec' for a CBOR based serialisation format, using strict
-- 'BS.ByteString's.
--
-- This is an adaptor between the @cborg@ library and the 'Codec' abstraction.
--
-- It takes encode and decode functions for the protocol messages that use the
-- CBOR library encoder and decoder.
--
-- Note that this is /less/ efficient than the 'mkCodecCborLazyBS' variant
-- because it has to copy and concatenate the result of the encoder (which
-- natively produces chunks).
--
mkCodecCborStrictBS
  :: forall ps f m. MonadST m

  => (forall (st :: ps) (st' :: ps).
             StateTokenI st
          =>ActiveState st
          => f st' -> Message ps st st' -> CBOR.Encoding)
  -- ^ cbor encoder

  -> (forall (st :: ps) s.
             ActiveState st
          => StateToken st
          -> f st
          -> CBOR.Decoder s (SomeMessage st))
  -- ^ cbor decoder

  -> Codec ps DeserialiseFailure f m BS.ByteString
mkCodecCborStrictBS :: forall ps (f :: ps -> *) (m :: * -> *).
MonadST m =>
(forall (st :: ps) (st' :: ps).
 (StateTokenI st, ActiveState st) =>
 f st' -> Message ps st st' -> Encoding)
-> (forall (st :: ps) s.
    ActiveState st =>
    StateToken st -> f st -> Decoder s (SomeMessage st))
-> Codec ps DeserialiseFailure f m ByteString
mkCodecCborStrictBS forall (st :: ps) (st' :: ps).
(StateTokenI st, ActiveState st) =>
f st' -> Message ps st st' -> Encoding
cborMsgEncode forall (st :: ps) s.
ActiveState st =>
StateToken st -> f st -> Decoder s (SomeMessage st)
cborMsgDecode =
    Codec {
      encode :: forall (st :: ps) (st' :: ps).
(StateTokenI st, ActiveState st) =>
f st' -> Message ps st st' -> ByteString
encode = \f st'
f Message ps st st'
msg  -> (Message ps st st' -> Encoding) -> Message ps st st' -> ByteString
forall a. (a -> Encoding) -> a -> ByteString
convertCborEncoder (f st' -> Message ps st st' -> Encoding
forall (st :: ps) (st' :: ps).
(StateTokenI st, ActiveState st) =>
f st' -> Message ps st st' -> Encoding
cborMsgEncode f st'
f) Message ps st st'
msg,
      decode :: forall (st :: ps).
ActiveState st =>
StateToken st
-> f st
-> m (DecodeStep ByteString DeserialiseFailure m (SomeMessage st))
decode = \StateToken st
stok f st
f -> (forall s. Decoder s (SomeMessage st))
-> m (DecodeStep ByteString DeserialiseFailure m (SomeMessage st))
forall a.
(forall s. Decoder s a)
-> m (DecodeStep ByteString DeserialiseFailure m a)
convertCborDecoder (StateToken st -> f st -> Decoder s (SomeMessage st)
forall (st :: ps) s.
ActiveState st =>
StateToken st -> f st -> Decoder s (SomeMessage st)
cborMsgDecode StateToken st
stok f st
f)
    }
  where
    convertCborEncoder :: (a -> CBOR.Encoding) -> a -> BS.ByteString
    convertCborEncoder :: forall a. (a -> Encoding) -> a -> ByteString
convertCborEncoder a -> Encoding
cborEncode =
        Encoding -> ByteString
CBOR.toStrictByteString
      (Encoding -> ByteString) -> (a -> Encoding) -> a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Encoding
cborEncode

    convertCborDecoder
      :: (forall s. CBOR.Decoder s a)
      -> m (DecodeStep BS.ByteString DeserialiseFailure m a)
    convertCborDecoder :: forall a.
(forall s. Decoder s a)
-> m (DecodeStep ByteString DeserialiseFailure m a)
convertCborDecoder forall s. Decoder s a
cborDecode =
        Decoder (PrimState m) a
-> (forall b. ST (PrimState m) b -> m b)
-> m (DecodeStep ByteString DeserialiseFailure m a)
forall s (m :: * -> *) a.
Functor m =>
Decoder s a
-> (forall b. ST s b -> m b)
-> m (DecodeStep ByteString DeserialiseFailure m a)
convertCborDecoderBS Decoder (PrimState m) a
forall s. Decoder s a
cborDecode ST (PrimState m) b -> m b
forall b. ST (PrimState m) b -> m b
forall (m :: * -> *) a. MonadST m => ST (PrimState m) a -> m a
stToIO

-- | Construct a 'Codec' for a CBOR based serialisation format, using lazy
-- 'BS.ByteString's.
--
-- This is an adaptor between the @cborg@ library and the 'Codec' abstraction.
--
-- It takes encode and decode functions for the protocol messages that use the
-- CBOR library encoder and decoder.
--
mkCodecCborLazyBS
  :: forall ps f m. MonadST m

  => (forall (st :: ps) (st' :: ps).
             StateTokenI st
          => ActiveState st
          => f st'
          -> Message ps st st' -> CBOR.Encoding)
  -- ^ cbor encoder

  -> (forall (st :: ps) s.
             ActiveState st
          => StateToken st
          -> f st
          -> CBOR.Decoder s (SomeMessage st))
  -- ^ cbor decoder

  -> Codec ps CBOR.DeserialiseFailure f m LBS.ByteString
mkCodecCborLazyBS :: forall ps (f :: ps -> *) (m :: * -> *).
MonadST m =>
(forall (st :: ps) (st' :: ps).
 (StateTokenI st, ActiveState st) =>
 f st' -> Message ps st st' -> Encoding)
-> (forall (st :: ps) s.
    ActiveState st =>
    StateToken st -> f st -> Decoder s (SomeMessage st))
-> Codec ps DeserialiseFailure f m ByteString
mkCodecCborLazyBS forall (st :: ps) (st' :: ps).
(StateTokenI st, ActiveState st) =>
f st' -> Message ps st st' -> Encoding
cborMsgEncode forall (st :: ps) s.
ActiveState st =>
StateToken st -> f st -> Decoder s (SomeMessage st)
cborMsgDecode =
    Codec {
      encode :: forall (st :: ps) (st' :: ps).
(StateTokenI st, ActiveState st) =>
f st' -> Message ps st st' -> ByteString
encode = \f st'
f Message ps st st'
msg  -> (Message ps st st' -> Encoding) -> Message ps st st' -> ByteString
forall a. (a -> Encoding) -> a -> ByteString
convertCborEncoder (f st' -> Message ps st st' -> Encoding
forall (st :: ps) (st' :: ps).
(StateTokenI st, ActiveState st) =>
f st' -> Message ps st st' -> Encoding
cborMsgEncode f st'
f) Message ps st st'
msg,
      decode :: forall (st :: ps).
ActiveState st =>
StateToken st
-> f st
-> m (DecodeStep ByteString DeserialiseFailure m (SomeMessage st))
decode = \StateToken st
stok f st
f -> (forall s. Decoder s (SomeMessage st))
-> m (DecodeStep ByteString DeserialiseFailure m (SomeMessage st))
forall a.
(forall s. Decoder s a)
-> m (DecodeStep ByteString DeserialiseFailure m a)
convertCborDecoder (StateToken st -> f st -> Decoder s (SomeMessage st)
forall (st :: ps) s.
ActiveState st =>
StateToken st -> f st -> Decoder s (SomeMessage st)
cborMsgDecode StateToken st
stok f st
f)
    }
  where
    convertCborEncoder :: (a -> CBOR.Encoding) -> a -> LBS.ByteString
    convertCborEncoder :: forall a. (a -> Encoding) -> a -> ByteString
convertCborEncoder a -> Encoding
cborEncode =
        Builder -> ByteString
toLazyByteString
      (Builder -> ByteString) -> (a -> Builder) -> a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Encoding -> Builder
CBOR.toBuilder
      (Encoding -> Builder) -> (a -> Encoding) -> a -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Encoding
cborEncode

    convertCborDecoder
      :: (forall s. CBOR.Decoder s a)
      -> m (DecodeStep LBS.ByteString CBOR.DeserialiseFailure m a)
    convertCborDecoder :: forall a.
(forall s. Decoder s a)
-> m (DecodeStep ByteString DeserialiseFailure m a)
convertCborDecoder forall s. Decoder s a
cborDecode =
        Decoder (PrimState m) a
-> (forall b. ST (PrimState m) b -> m b)
-> m (DecodeStep ByteString DeserialiseFailure m a)
forall s (m :: * -> *) a.
Monad m =>
Decoder s a
-> (forall b. ST s b -> m b)
-> m (DecodeStep ByteString DeserialiseFailure m a)
convertCborDecoderLBS Decoder (PrimState m) a
forall s. Decoder s a
cborDecode ST (PrimState m) b -> m b
forall b. ST (PrimState m) b -> m b
forall (m :: * -> *) a. MonadST m => ST (PrimState m) a -> m a
stToIO

{-# NOINLINE toLazyByteString #-}
toLazyByteString :: BS.Builder -> LBS.ByteString
toLazyByteString :: Builder -> ByteString
toLazyByteString = AllocationStrategy -> ByteString -> Builder -> ByteString
BS.toLazyByteStringWith AllocationStrategy
strategy ByteString
LBS.empty
  where
    -- Buffer strategy and sizes better tuned to our network protocol situation.
    --
    -- The LBS.smallChunkSize is 4k - heap object overheads, so that
    -- it does fit in a 4k overall.
    --
    strategy :: AllocationStrategy
strategy = Int -> Int -> AllocationStrategy
BS.untrimmedStrategy Int
800 Int
LBS.smallChunkSize