module Data.Packed.TH.Write (genWrite, writeFName) where

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

-- For a data type 'Tree', will generate the function name 'writeTree'
writeFName :: Name -> Name
writeFName :: Name -> Name
writeFName Name
tyName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"write" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
sanitizeConName Name
tyName

-- | Generates a function that serialises and writes a value into a 'Needs'
--
-- The function simply calls the functions generated by 'Data.Packed.TH.genConWrite'
--
-- __Example:__
--
-- For the 'Tree' data type, it generates the following function
--
-- @
-- writeTree :: ('Packable' a) => Tree a -> 'NeedsWriter' (Tree a) r t
-- writeTree (Leaf n) = writeConLeaf n
-- writeTree (Node l r) = writeConNode l r
-- @
genWrite ::
    [PackingFlag] ->
    -- | The name of the type to generate the function for
    Name ->
    Q [Dec]
genWrite :: [PackingFlag] -> Name -> Q [Dec]
genWrite [PackingFlag]
flags Name
tyName = do
    (TyConI (DataD _ _ _ _ cs _)) <- Name -> Q Info
reify Name
tyName
    -- For each data constructor, we generate the corresponding clause
    clauses <-
        mapM
            ( \Con
con -> do
                let (Name
conName, [BangType]
types) = Con -> (Name, [BangType])
getNameAndBangTypesFromCon Con
con
                -- Generate names for each variable in the constructor
                paramNames <- (BangType -> Q Name) -> [BangType] -> 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 (Q Name -> BangType -> Q Name
forall a b. a -> b -> a
const (Q Name -> BangType -> Q Name) -> Q Name -> BangType -> Q Name
forall a b. (a -> b) -> a -> b
$ String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"t") [BangType]
types
                -- We apply each parameter of the constructor and the 'Needs' to the 'writeConXXX' function
                body <- foldl (\Q Exp
f Name
arg -> [|$Q Exp
f $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
arg)|]) (varE $ conWriteFName conName) paramNames
                return $ Clause [ConP conName [] (VarP <$> paramNames)] (NormalB body) []
            )
            cs
    -- For each of the data constructor of the type, we generate the corresponding `writeConXXX`
    -- We define the Tag using the index of the data constructor
    conWriter <- forM (zip [0 ..] cs) (\(Tag
index, Con
con) -> [PackingFlag] -> Con -> Tag -> Q [Dec]
genConWrite [PackingFlag]
flags Con
con Tag
index)
    signature <- genWriteSignature tyName
    return $ concat conWriter ++ [signature, FunD (writeFName tyName) clauses]

-- Generates the following function signature for a data type 'Tree'
-- writeTree :: ('Packable' a) => Tree a -> 'NeedsWriter' (Tree a) r t
genWriteSignature :: Name -> Q Dec
genWriteSignature :: Name -> Q Dec
genWriteSignature Name
tyName = do
    (sourceType, typeParameterNames) <- Name -> Q (Kind, [Name])
resolveAppliedType Name
tyName
    let fName = Name -> Name
writeFName Name
tyName
        -- Type variables for Needs
        r = Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT (Name -> Q Kind) -> Name -> Q Kind
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"r"
        t = Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT (Name -> Q Kind) -> Name -> Q Kind
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"t"
        -- Define Packable constraints on each of the type parameters
        constraints = (Name -> Q Kind) -> [Name] -> Q Cxt
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 (\Name
tyVarName -> [t|Packable $(Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT Name
tyVarName)|]) [Name]
typeParameterNames
        signature = [t|$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
sourceType) -> NeedsWriter $(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
sourceType) $Q Kind
r $Q Kind
t|]
    sigD fName (forallT [] constraints signature)