module Data.Packed.TH.RepackCon (genConstructorRepackers) where

import Data.Packed.FieldSize
import Data.Packed.Needs (Needs, applyNeeds, withEmptyNeeds)
import qualified Data.Packed.Needs as N
import Data.Packed.TH.Flag
import Data.Packed.TH.Start (startFName)
import Data.Packed.TH.Utils
import Language.Haskell.TH

-- | Generates a function that builds back data using already serialised fields
--
-- __Example:__
--
-- For the 'Tree' data type, it generates the following functions
--
-- @
-- repackLeaf :: 'Data.Packed.Needs' '[] a -> 'Data.Packed.Needs' '[] (Tree a)
-- repackLeaf pval = withEmptyNeeds (startLeaf N.>> 'Data.Packed.Needs.concatNeeds' pval)
--
-- repackNode :: 'Data.Packed.Needs' '[] (Tree a) -> 'Data.Packed.Needs' '[] (Tree a) -> 'Data.Packed.Needs' '[] (Tree a)
-- repackNode lval rval = withEmptyNeeds (startNode N.>> 'concatNeeds' lval N.>> 'concatNeeds' rval)
-- @
genConstructorRepackers :: [PackingFlag] -> Name -> Q [Dec]
genConstructorRepackers :: [PackingFlag] -> Name -> Q [Dec]
genConstructorRepackers [PackingFlag]
flags Name
tyName = do
    (TyConI (DataD _ _ _ _ cs _)) <- Name -> Q Info
reify Name
tyName
    packers <- genConstructorRepacker flags `mapM` cs
    return $ concat packers

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

genConstructorRepacker :: [PackingFlag] -> Con -> Q [Dec]
genConstructorRepacker :: [PackingFlag] -> Con -> Q [Dec]
genConstructorRepacker [PackingFlag]
flags Con
con = do
    let conName :: Name
conName = (Name, [BangType]) -> Name
forall a b. (a, b) -> a
fst ((Name, [BangType]) -> Name) -> (Name, [BangType]) -> Name
forall a b. (a -> b) -> a -> b
$ Con -> (Name, [BangType])
getNameAndBangTypesFromCon Con
con
        fieldTypes :: [(Kind, Int, Bool)]
fieldTypes = Con -> [PackingFlag] -> [(Kind, Int, Bool)]
getConFieldsIdxAndNeedsFS Con
con [PackingFlag]
flags
    varNames <- ((Kind, Int, Bool) -> Q Name) -> [(Kind, Int, Bool)] -> 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, Int, Bool)
_ -> String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"t") [(Kind, Int, Bool)]
fieldTypes
    writeExp <-
        let concated =
                (Q Exp -> ((Kind, Int, Bool), Name) -> Q Exp)
-> Q Exp -> [((Kind, Int, Bool), Name)] -> Q Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
                    ( \Q Exp
rest ((Kind
_, Int
_, Bool
needsFieldSize), Name
varName) ->
                        if Bool
needsFieldSize
                            then [|($Q Exp
rest) N.>> applyNeedsWithFieldSize $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
varName)|]
                            else [|($Q Exp
rest) N.>> applyNeeds $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
varName)|]
                    )
                    [|$(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Name -> Q Exp) -> Name -> Q Exp
forall a b. (a -> b) -> a -> b
$ Name -> Name
startFName Name
conName)|]
                    ([(Kind, Int, Bool)] -> [Name] -> [((Kind, Int, Bool), Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Kind, Int, Bool)]
fieldTypes [Name]
varNames)
         in [|withEmptyNeeds $concated|]
    signature <- genConstructorPackerSig flags conName ((\(Kind
t, Int
_, Bool
_) -> Kind
t) <$> fieldTypes)
    return
        [ signature
        , FunD (repackConFName conName) [Clause (VarP <$> varNames) (NormalB writeExp) []]
        ]

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, _) <- resolveAppliedType tyName
    signature <- foldr (\Kind
p Q Kind
rest -> [t|Needs '[] '[$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
p)] -> $Q Kind
rest|]) [t|Needs '[] '[$(return ty)]|] argTypes
    return $ SigD (repackConFName conName) $ ForallT [] [] signature