module Data.Packed.TH.PackCon (genConstructorPackers) where

import Data.Packed.Needs
import Data.Packed.Packable
import Data.Packed.Packed
import Data.Packed.TH.Flag (PackingFlag)
import Data.Packed.TH.Utils
import Data.Packed.TH.WriteCon
import Language.Haskell.TH

-- | Generates a function that serialises an applied data constructor
--
-- The function calls the functions generated by 'Data.Packed.TH.genConWrite'
--
-- __Example:__
--
-- For the 'Tree' data type, it generates the following functions
--
-- @
-- packLeaf :: ('Packable' a) => a -> 'Data.Packed' '[Tree a]
-- packLeaf n = 'finish' ('withEmptyNeeds' (writeLeaf n))
--
-- packNode :: ('Packable' a) => Tree a -> Tree a -> 'Data.Packed' '[Tree a]
-- packNode t1 t2 = 'finish' ('withEmptyNeeds' (writeNode t1 t2))
-- @
genConstructorPackers :: [PackingFlag] -> Name -> Q [Dec]
genConstructorPackers :: [PackingFlag] -> Name -> Q [Dec]
genConstructorPackers [PackingFlag]
flags Name
tyName = do
    (TyConI (DataD _ _ _ _ cs _)) <- Name -> Q Info
reify Name
tyName
    packers <-
        mapM
            ( \Con
con ->
                let (Name
conName, [BangType]
bt) = Con -> (Name, [BangType])
getNameAndBangTypesFromCon Con
con
                 in [PackingFlag] -> Name -> Cxt -> Q [Dec]
genConstructorPacker [PackingFlag]
flags Name
conName (BangType -> Kind
forall a b. (a, b) -> b
snd (BangType -> Kind) -> [BangType] -> Cxt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [BangType]
bt)
            )
            cs
    return $ concat packers

packConFName :: Name -> Name
packConFName :: Name -> Name
packConFName Name
conName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"pack" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
sanitizeConName Name
conName

genConstructorPacker :: [PackingFlag] -> Name -> [Type] -> Q [Dec]
genConstructorPacker :: [PackingFlag] -> Name -> Cxt -> Q [Dec]
genConstructorPacker [PackingFlag]
flags Name
conName Cxt
argTypes = do
    varNames <- (Kind -> Q Name) -> Cxt -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Kind
_ -> String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"t") Cxt
argTypes
    writeExp <- (foldl (\Q Exp
rest Name
p -> Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE Q Exp
rest (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
p)) (varE $ conWriteFName conName) varNames)
    body <- [|finish (withEmptyNeeds $(return writeExp))|]
    signature <- genConstructorPackerSig flags conName argTypes
    return
        [ signature
        , FunD (packConFName conName) [Clause (VarP <$> varNames) (NormalB body) []]
        ]

genConstructorPackerSig :: [PackingFlag] -> Name -> [Type] -> Q Dec
genConstructorPackerSig :: [PackingFlag] -> Name -> Cxt -> Q Dec
genConstructorPackerSig [PackingFlag]
_ Name
conName Cxt
argTypes = do
    (DataConI _ _ tyName) <- Name -> Q Info
reify Name
conName
    (ty, typeParameterNames) <- resolveAppliedType tyName
    constraints <- mapM (\Name
tyVarName -> [t|Packable $(Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT Name
tyVarName)|]) typeParameterNames
    signature <- foldr (\Kind
p Q Kind
rest -> [t|$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
p) -> $Q Kind
rest|]) [t|Packed '[$(return ty)]|] argTypes
    return $ SigD (packConFName conName) $ ForallT [] constraints signature