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 <-
        mapM
            ( \Con
con ->
                let (Name
conName, [BangType]
bt) = Con -> (Name, [BangType])
getNameAndBangTypesFromCon Con
con
                 in [PackingFlag] -> Name -> Cxt -> Q [Dec]
genConstructorRepacker [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

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] -> Name -> [Type] -> Q [Dec]
genConstructorRepacker :: [PackingFlag] -> Name -> Cxt -> Q [Dec]
genConstructorRepacker [PackingFlag]
flags Name
conName Cxt
argTypes = do
    let argCount :: Int
argCount = Cxt -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Cxt
argTypes
        needsFieldSize :: Int -> Bool
needsFieldSize Int
i =
            (PackingFlag
InsertFieldSize PackingFlag -> [PackingFlag] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [PackingFlag]
flags)
                Bool -> Bool -> Bool
&& ( (PackingFlag
SkipLastFieldSize PackingFlag -> [PackingFlag] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [PackingFlag]
flags)
                        Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
argCount Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                   )
    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 <-
        let concated =
                (Q Exp -> (Int, Name) -> Q Exp) -> Q Exp -> [(Int, 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 (Int
i, Name
p) ->
                        if Int -> Bool
needsFieldSize Int
i
                            then [|($Q Exp
rest) N.>> applyNeedsWithFieldSize $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
p)|]
                            else [|($Q Exp
rest) N.>> applyNeeds $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
p)|]
                    )
                    [|$(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)|]
                    ([Int] -> [Name] -> [(Int, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Name]
varNames)
         in [|withEmptyNeeds $concated|]
    signature <- genConstructorPackerSig flags conName argTypes
    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