Generic.hs

{-# LANGUAGE UndecidableInstances #-}

-- | Internal module. Not part of the public API.
module Data.PackStream.Generic
    ( GPackStream
    , genericToPs
    , genericFromPs
    , GenericPackStream(..)
    ) where

import           Compat.Prelude

import           Data.Kind (Constraint, Type)
import           GHC.Generics

import           Data.PackStream.Ps
import           Data.PackStream.Result

-- | Convert a value with a 'Generic' instance to a 'Ps' value.
genericToPs :: (Generic a, GPackStream (Rep a)) => a -> Ps
genericToPs = gToPs . from

-- | Decode a 'Ps' value into a type with a 'Generic' instance.
genericFromPs :: (Generic a, GPackStream (Rep a)) => Ps -> Result a
genericFromPs x = to <$> gFromPs x

-- | Newtype wrapper providing a 'PackStream' instance via 'GHC.Generics'.
type GenericPackStream :: Type -> Type
type role GenericPackStream representational
newtype GenericPackStream a = GenericPackStream a

instance (Generic a, GPackStream (Rep a)) => PackStream (GenericPackStream a) where
  toPs (GenericPackStream a) = genericToPs a
  fromPs a = GenericPackStream <$> genericFromPs a

-- | Generic representation class for PackStream encoding\/decoding.
type GPackStream :: (Type -> Type) -> Constraint
class GPackStream f where
  gToPs   :: f a -> Ps
  gFromPs :: Ps -> Result (f a)

instance GPackStream U1 where
  gToPs U1 = PsNull
  gFromPs PsNull = return U1
  gFromPs _      = fail "invalid encoding for custom unit type"

instance (GPackStream a, GProdPack b) => GPackStream (a :*: b) where
  gToPs = toPs . prodToPs
  gFromPs = fromPs >=> prodFromPs

instance (GSumPack a, GSumPack b, SumSize a, SumSize b) => GPackStream (a :+: b) where
  gToPs = sumToPs 0 size
    where
      size = unTagged (sumSize :: Tagged (a :+: b) Int64)

  gFromPs = \case
    PsInteger code -> checkSumFromPs0 size (fromIntegral code)
    o              -> fromPs o >>= uncurry (checkSumFromPs size)
    where
      size = unTagged (sumSize :: Tagged (a :+: b) Int64)

instance GPackStream a => GPackStream (M1 t c a) where
  gToPs (M1 x) = gToPs x
  gFromPs x = M1 <$> gFromPs x

instance PackStream a => GPackStream (K1 i a) where
  gToPs (K1 x) = toPs x
  gFromPs o = K1 <$> fromPs o


-- Product type packing.

type GProdPack :: (Type -> Type) -> Constraint
class GProdPack f where
  prodToPs :: f a -> [Ps]
  prodFromPs :: [Ps] -> Result (f a)


instance (GPackStream a, GProdPack b) => GProdPack (a :*: b) where
  prodToPs (a :*: b) = gToPs a : prodToPs b
  prodFromPs (a:b) = (:*:) <$> gFromPs a <*> prodFromPs b
  prodFromPs _     = fail "invalid encoding for product type"

instance GPackStream a => GProdPack (M1 t c a) where
  prodToPs (M1 x) = [gToPs x]
  prodFromPs [x] = M1 <$> gFromPs x
  prodFromPs _   = fail "invalid encoding for product type"


-- Sum type packing.

checkSumFromPs0 :: GSumPack f => Int64 -> Int64 -> Result (f a)
checkSumFromPs0 size code
  | code < size = sumFromPs code size PsNull
  | otherwise   = fail "invalid encoding for sum type"


checkSumFromPs :: (GSumPack f) => Int64 -> Int64 -> Ps -> Result (f a)
checkSumFromPs size code x
  | code < size = sumFromPs code size x
  | otherwise   = fail "invalid encoding for sum type"


type GSumPack :: (Type -> Type) -> Constraint
class GSumPack f where
  sumToPs :: Int64 -> Int64 -> f a -> Ps
  sumFromPs :: Int64 -> Int64 -> Ps -> Result (f a)


instance (GSumPack a, GSumPack b) => GSumPack (a :+: b) where
  sumToPs code size = \case
    L1 x -> sumToPs code           sizeL x
    R1 x -> sumToPs (code + sizeL) sizeR x
    where
      sizeL = size `shiftR` 1
      sizeR = size - sizeL

  sumFromPs code size x
    | code < sizeL = L1 <$> sumFromPs code           sizeL x
    | otherwise    = R1 <$> sumFromPs (code - sizeL) sizeR x
    where
      sizeL = size `shiftR` 1
      sizeR = size - sizeL


instance {-# OVERLAPPING #-} GSumPack (C1 c U1) where
  sumToPs code _ _ = toPs code
  sumFromPs _ _ = gFromPs


instance {-# OVERLAPPABLE #-} GPackStream a => GSumPack (C1 c a) where
  sumToPs code _ x = toPs (code, gToPs x)
  sumFromPs _ _ = gFromPs


-- Sum size.

type SumSize :: (Type -> Type) -> Constraint
class SumSize f where
  sumSize :: Tagged f Int64

type Tagged :: (Type -> Type) -> Type -> Type
type role Tagged phantom representational
newtype Tagged (s :: Type -> Type) b = Tagged { unTagged :: b }

instance (SumSize a, SumSize b) => SumSize (a :+: b) where
  sumSize = Tagged $ unTagged (sumSize :: Tagged a Int64) +
                     unTagged (sumSize :: Tagged b Int64)

instance SumSize (C1 c a) where
  sumSize = Tagged 1