{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}

module Data.ProtocolBuffers.Encode
  ( Encode(..)
  , encodeMessage
  , encodeLengthPrefixedMessage
  , GEncode
  ) where

import qualified Data.ByteString as B
import Data.Foldable
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.Proxy
import Data.Serialize.Put

import GHC.Generics
import GHC.TypeLits

import Data.ProtocolBuffers.Types
import Data.ProtocolBuffers.Wire

-- |
-- Encode a Protocol Buffers message.
encodeMessage :: Encode a => a -> Put
encodeMessage :: forall a. Encode a => a -> Put
encodeMessage = a -> Put
forall a. Encode a => a -> Put
encode

-- |
-- Encode a Protocol Buffers message prefixed with a varint encoded 32-bit integer describing its length.
encodeLengthPrefixedMessage :: Encode a => a -> Put
{-# INLINE encodeLengthPrefixedMessage #-}
encodeLengthPrefixedMessage :: forall a. Encode a => a -> Put
encodeLengthPrefixedMessage a
msg = do
  let msg' :: ByteString
msg' = Put -> ByteString
runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ a -> Put
forall a. Encode a => a -> Put
encodeMessage a
msg
  Int -> Put
forall a. (Integral a, Bits a) => a -> Put
putVarUInt (Int -> Put) -> Int -> Put
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
B.length ByteString
msg'
  Putter ByteString
putByteString ByteString
msg'

class Encode (a :: *) where
  encode :: a -> Put
  default encode :: (Generic a, GEncode (Rep a)) => a -> Put
  encode = Rep a Any -> Put
forall a. Rep a a -> Put
forall (f :: * -> *) a. GEncode f => f a -> Put
gencode (Rep a Any -> Put) -> (a -> Rep a Any) -> a -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Rep a Any
forall x. a -> Rep a x
forall a x. Generic a => a -> Rep a x
from

-- | Untyped message encoding
instance Encode (HashMap Tag [WireField]) where
  encode :: HashMap Tag [WireField] -> Put
encode = ((Tag, [WireField]) -> Put) -> [(Tag, [WireField])] -> Put
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Tag, [WireField]) -> Put
step ([(Tag, [WireField])] -> Put)
-> (HashMap Tag [WireField] -> [(Tag, [WireField])])
-> HashMap Tag [WireField]
-> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap Tag [WireField] -> [(Tag, [WireField])]
forall k v. HashMap k v -> [(k, v)]
HashMap.toList where
    step :: (Tag, [WireField]) -> Put
step = (Tag -> [WireField] -> Put) -> (Tag, [WireField]) -> Put
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((WireField -> Put) -> [WireField] -> Put
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ ((WireField -> Put) -> [WireField] -> Put)
-> (Tag -> WireField -> Put) -> Tag -> [WireField] -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tag -> WireField -> Put
forall a. EncodeWire a => Tag -> a -> Put
encodeWire)

class GEncode (f :: * -> *) where
  gencode :: f a -> Put

instance GEncode a => GEncode (M1 i c a) where
  gencode :: forall a. M1 i c a a -> Put
gencode = a a -> Put
forall a. a a -> Put
forall (f :: * -> *) a. GEncode f => f a -> Put
gencode (a a -> Put) -> (M1 i c a a -> a a) -> M1 i c a a -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. M1 i c a a -> a a
forall k i (c :: Meta) (f :: k -> *) (p :: k). M1 i c f p -> f p
unM1

instance (GEncode a, GEncode b) => GEncode (a :*: b) where
  gencode :: forall a. (:*:) a b a -> Put
gencode (a a
x :*: b a
y) = a a -> Put
forall a. a a -> Put
forall (f :: * -> *) a. GEncode f => f a -> Put
gencode a a
x Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> b a -> Put
forall a. b a -> Put
forall (f :: * -> *) a. GEncode f => f a -> Put
gencode b a
y

instance (GEncode a, GEncode b) => GEncode (a :+: b) where
  gencode :: forall a. (:+:) a b a -> Put
gencode (L1 a a
x) = a a -> Put
forall a. a a -> Put
forall (f :: * -> *) a. GEncode f => f a -> Put
gencode a a
x
  gencode (R1 b a
y) = b a -> Put
forall a. b a -> Put
forall (f :: * -> *) a. GEncode f => f a -> Put
gencode b a
y

instance (EncodeWire a, KnownNat n, Foldable f) => GEncode (K1 i (Field n (f a))) where
  gencode :: forall a. K1 i (Field n (f a)) a -> Put
gencode = (a -> Put) -> f a -> Put
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Tag -> a -> Put
forall a. EncodeWire a => Tag -> a -> Put
encodeWire Tag
tag) (f a -> Put)
-> (K1 i (Field n (f a)) a -> f a) -> K1 i (Field n (f a)) a -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field n (f a) -> f a
forall (n :: Nat) a. Field n a -> a
runField (Field n (f a) -> f a)
-> (K1 i (Field n (f a)) a -> Field n (f a))
-> K1 i (Field n (f a)) a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. K1 i (Field n (f a)) a -> Field n (f a)
forall k i c (p :: k). K1 i c p -> c
unK1 where
    tag :: Tag
tag = Integer -> Tag
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Integer -> Tag) -> Integer -> Tag
forall a b. (a -> b) -> a -> b
$ Proxy n -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy n
forall {k} (t :: k). Proxy t
Proxy :: Proxy n)

instance GEncode U1 where
  gencode :: forall a. U1 a -> Put
gencode U1 a
_ = () -> Put
forall a. a -> PutM a
forall (m :: * -> *) a. Monad m => a -> m a
return ()