{-# LANGUAGE BangPatterns      #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes        #-}
{-# LANGUAGE TypeFamilies      #-}
{-# LANGUAGE TypeInType        #-}
{-# LANGUAGE TypeOperators     #-}

-- |
-- Module      :  Codec.Scale.Generic
-- Copyright   :  Aleksandr Krupenkin 2016-2021
-- License     :  Apache-2.0
--
-- Maintainer  :  mail@akru.me
-- Stability   :  experimental
-- Portability :  unportable
--
-- This module defines generic codec instances for data structures (including tuples)
-- and enums (tagged-unions in Rust).
--

module Codec.Scale.Generic () where

import           Data.Serialize.Get (Get, getWord8)
import           Data.Serialize.Put (PutM, putWord8)
import           Data.Word          (Word8)
import           Generics.SOP       (All, Compose, I (..), NP (..), NS (..),
                                     SOP (..), unSOP, unZ)

import           Codec.Scale.Class  (Decode (..), Encode (..), GDecode (..),
                                     GEncode (..))

-- Enum has multiple sum types.
instance ( GEncode (NP f xs)
         , GEncode (NP f ys)
         , All (GEncode `Compose` NP f) xss
         ) => GEncode (SOP f (xs ': ys ': xss)) where
    gPut :: Putter (SOP f (xs : ys : xss))
gPut = Word8 -> NS (NP f) (xs : ys : xss) -> PutM ()
forall k (f :: k -> *) (as :: [k]).
All (Compose GEncode f) as =>
Word8 -> NS f as -> PutM ()
go Word8
0 (NS (NP f) (xs : ys : xss) -> PutM ())
-> (SOP f (xs : ys : xss) -> NS (NP f) (xs : ys : xss))
-> Putter (SOP f (xs : ys : xss))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOP f (xs : ys : xss) -> NS (NP f) (xs : ys : xss)
forall k (f :: k -> *) (xss :: [[k]]). SOP f xss -> NS (NP f) xss
unSOP
      where
        go :: forall f as . All (GEncode `Compose` f) as => Word8 -> NS f as -> PutM ()
        go :: Word8 -> NS f as -> PutM ()
go !Word8
acc (Z f x
x) = Putter Word8
putWord8 Word8
acc PutM () -> PutM () -> PutM ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Putter (f x)
forall a. GEncode a => Putter a
gPut f x
x
        go !Word8
acc (S NS f xs
x) = Word8 -> NS f xs -> PutM ()
forall k (f :: k -> *) (as :: [k]).
All (Compose GEncode f) as =>
Word8 -> NS f as -> PutM ()
go (Word8
acc Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word8
1) NS f xs
x

-- Structures has only one sum type.
instance GEncode (NP f xs) => GEncode (SOP f '[xs]) where
    gPut :: Putter (SOP f '[xs])
gPut = Putter (NP f xs)
forall a. GEncode a => Putter a
gPut Putter (NP f xs)
-> (SOP f '[xs] -> NP f xs) -> Putter (SOP f '[xs])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NS (NP f) '[xs] -> NP f xs
forall k (f :: k -> *) (x :: k). NS f '[x] -> f x
unZ (NS (NP f) '[xs] -> NP f xs)
-> (SOP f '[xs] -> NS (NP f) '[xs]) -> SOP f '[xs] -> NP f xs
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SOP f '[xs] -> NS (NP f) '[xs]
forall k (f :: k -> *) (xss :: [[k]]). SOP f xss -> NS (NP f) xss
unSOP

-- Product serialization is just encode each field step by step.
instance (Encode a, GEncode (NP I as)) => GEncode (NP I (a ': as)) where
    gPut :: Putter (NP I (a : as))
gPut (I x
a :* NP I xs
as) = Putter x
forall a. Encode a => Putter a
put x
a PutM () -> PutM () -> PutM ()
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Putter (NP I xs)
forall a. GEncode a => Putter a
gPut NP I xs
as

-- Finish when all fields handled.
instance GEncode (NP I '[]) where
    gPut :: Putter (NP I '[])
gPut NP I '[]
_ = PutM ()
forall a. Monoid a => a
mempty

-- | Enum parser definition.
--
-- The index of sum type to parse given as an argument.
class EnumParser xs where
    enumParser :: All (GDecode `Compose` NP f) xs => Word8 -> Get (NS (NP f) xs)

-- Enumerate enum index, zero means that we reach the goal.
instance EnumParser as => EnumParser (a ': as) where
    enumParser :: Word8 -> Get (NS (NP f) (a : as))
enumParser !Word8
i | Word8
i Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
> Word8
0     = NS (NP f) as -> NS (NP f) (a : as)
forall k (a :: k -> *) (xs :: [k]) (x :: k).
NS a xs -> NS a (x : xs)
S (NS (NP f) as -> NS (NP f) (a : as))
-> Get (NS (NP f) as) -> Get (NS (NP f) (a : as))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Word8 -> Get (NS (NP f) as)
forall k (xs :: [[k]]) (f :: k -> *).
(EnumParser xs, All (Compose GDecode (NP f)) xs) =>
Word8 -> Get (NS (NP f) xs)
enumParser (Word8
i Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
- Word8
1)
                  | Bool
otherwise = NP f a -> NS (NP f) (a : as)
forall k (a :: k -> *) (x :: k) (xs :: [k]). a x -> NS a (x : xs)
Z (NP f a -> NS (NP f) (a : as))
-> Get (NP f a) -> Get (NS (NP f) (a : as))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get (NP f a)
forall a. GDecode a => Get a
gGet

-- When index out of type scope raise the error.
instance EnumParser '[] where
    enumParser :: Word8 -> Get (NS (NP f) '[])
enumParser Word8
i = String -> Get (NS (NP f) '[])
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"index out of enum constructors count: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Word8 -> String
forall a. Show a => a -> String
show Word8
i)

-- Decode enum when multiple sum types.
instance ( GDecode (NP f xs)
         , GDecode (NP f ys)
         , All (GDecode `Compose` NP f) xss
         , EnumParser xss
         ) => GDecode (SOP f (xs ': ys ': xss)) where
    gGet :: Get (SOP f (xs : ys : xss))
gGet = NS (NP f) (xs : ys : xss) -> SOP f (xs : ys : xss)
forall k (f :: k -> *) (xss :: [[k]]). NS (NP f) xss -> SOP f xss
SOP (NS (NP f) (xs : ys : xss) -> SOP f (xs : ys : xss))
-> Get (NS (NP f) (xs : ys : xss)) -> Get (SOP f (xs : ys : xss))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Word8 -> Get (NS (NP f) (xs : ys : xss))
forall k (xs :: [[k]]) (f :: k -> *).
(EnumParser xs, All (Compose GDecode (NP f)) xs) =>
Word8 -> Get (NS (NP f) xs)
enumParser (Word8 -> Get (NS (NP f) (xs : ys : xss)))
-> Get Word8 -> Get (NS (NP f) (xs : ys : xss))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Get Word8
getWord8)

-- Decode plain structure when only one sum type.
instance GDecode (NP f as) => GDecode (SOP f '[as]) where
    gGet :: Get (SOP f '[as])
gGet = NS (NP f) '[as] -> SOP f '[as]
forall k (f :: k -> *) (xss :: [[k]]). NS (NP f) xss -> SOP f xss
SOP (NS (NP f) '[as] -> SOP f '[as])
-> (NP f as -> NS (NP f) '[as]) -> NP f as -> SOP f '[as]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NP f as -> NS (NP f) '[as]
forall k (a :: k -> *) (x :: k) (xs :: [k]). a x -> NS a (x : xs)
Z (NP f as -> SOP f '[as]) -> Get (NP f as) -> Get (SOP f '[as])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get (NP f as)
forall a. GDecode a => Get a
gGet

-- Decode each field in sequence.
instance (Decode a, GDecode (NP I as)) => GDecode (NP I (a ': as)) where
    gGet :: Get (NP I (a : as))
gGet = I a -> NP I as -> NP I (a : as)
forall k (a :: k -> *) (x :: k) (xs :: [k]).
a x -> NP a xs -> NP a (x : xs)
(:*) (I a -> NP I as -> NP I (a : as))
-> Get (I a) -> Get (NP I as -> NP I (a : as))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> I a
forall a. a -> I a
I (a -> I a) -> Get a -> Get (I a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get a
forall a. Decode a => Get a
get) Get (NP I as -> NP I (a : as))
-> Get (NP I as) -> Get (NP I (a : as))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get (NP I as)
forall a. GDecode a => Get a
gGet

-- Finish decoding when empty.
instance GDecode (NP I '[]) where
    gGet :: Get (NP I '[])
gGet = NP I '[] -> Get (NP I '[])
forall (m :: * -> *) a. Monad m => a -> m a
return NP I '[]
forall k (a :: k -> *). NP a '[]
Nil