module Data.Packed.TH.Transform (transformFName, genTransform) where

import Data.Maybe (catMaybes)
import Data.Packed.FieldSize (FieldSize)
import Data.Packed.Needs (withEmptyNeeds)
import qualified Data.Packed.Needs as N
import Data.Packed.Reader (PackedReader)
import qualified Data.Packed.Reader as R
import Data.Packed.TH.Case (caseFName)
import Data.Packed.TH.Flag
import Data.Packed.TH.Start (startFName)
import Data.Packed.TH.Utils
import Language.Haskell.TH

-- | For a constructor 'Leaf', will generate the function name 'transformLeaf'
transformFName :: Name -> Name
transformFName :: Name -> Name
transformFName Name
conName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"transform" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
sanitizeConName Name
conName

-- For a type 'Tree', generates the following function
--
-- transformTree ::
--     ('Data.Packed.Reader.PackedReader' '[a] r ('Data.Packed.Needs.NeedsBuilder' '[a] '[Tree a] '[] '[Tree a])) ->
--
--     ('Data.Packed.Reader.PackedReader' '[Tree a, Tree a] r ('Data.Packed.Needs.NeedsBuilder' '[Tree a, Tree a] '[Tree a] '[] '[Tree a])) ->
--     'Data.Packed.PackedReader' '[Tree a] r ('Data.Packed.Needs' '[] '[Tree a])
-- transformTree leafCase nodeCase = caseTree
--      (leafCase R.>>= \l -> 'Data.Packed.Needs.finish' ('Data.Packed.Needs.withEmptyNeeds' (startLeaf 'Data.Packed.Needs.>>' l)))
--      (nodeCase R.>>= \n -> 'Data.Packed.Needs.finish' ('Data.Packed.Needs.withEmptyNeeds' (startNode 'Data.Packed.Needs.>>' n)))
genTransform :: [PackingFlag] -> Name -> Q [Dec]
genTransform :: [PackingFlag] -> Name -> Q [Dec]
genTransform [PackingFlag]
flags Name
tyName = do
    signature <- [PackingFlag] -> Name -> Q Dec
genTransformSignature [PackingFlag]
flags Name
tyName
    (TyConI (DataD _ _ _ _ cs _)) <- reify tyName
    body <-
        foldl
            ( \Q Exp
rest Con
curr ->
                let caseName :: Name
caseName = Con -> Name
buildCaseFunctionName Con
curr
                 in if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Con -> Bool
conHasArguments Con
curr
                        then [|$Q Exp
rest (R.return (withEmptyNeeds $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Con -> Name
startFNameForCon Con
curr))))|]
                        else [|$Q Exp
rest ($(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
caseName) R.>>= \resWriter -> R.return (withEmptyNeeds ($(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Con -> Name
startFNameForCon Con
curr)) N.>> resWriter)))|]
            )
            (varE $ caseFNameForType tyName)
            cs
    return
        [ signature
        , FunD
            (transformFName tyName)
            [Clause (VarP . buildCaseFunctionName <$> filter conHasArguments cs) (NormalB body) []]
        ]
  where
    -- for dataconstructor Leaf, will be 'leafCase'
    buildCaseFunctionName :: Con -> Name
buildCaseFunctionName = Name -> Name
conNameToCaseFunctionName (Name -> Name) -> (Con -> Name) -> Con -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [BangType]) -> Name
forall a b. (a, b) -> a
fst ((Name, [BangType]) -> Name)
-> (Con -> (Name, [BangType])) -> Con -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> (Name, [BangType])
getNameAndBangTypesFromCon
    conNameToCaseFunctionName :: Name -> Name
conNameToCaseFunctionName Name
conName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"case" String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Name -> String
sanitizeConName Name
conName)

    startFNameForCon :: Con -> Name
startFNameForCon = Name -> Name
startFName (Name -> Name) -> (Con -> Name) -> Con -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [BangType]) -> Name
forall a b. (a, b) -> a
fst ((Name, [BangType]) -> Name)
-> (Con -> (Name, [BangType])) -> Con -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> (Name, [BangType])
getNameAndBangTypesFromCon
    caseFNameForType :: Name -> Name
caseFNameForType = Name -> Name
caseFName
    conHasArguments :: Con -> Bool
conHasArguments = Bool -> Bool
not (Bool -> Bool) -> (Con -> Bool) -> Con -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [BangType] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([BangType] -> Bool) -> (Con -> [BangType]) -> Con -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [BangType]) -> [BangType]
forall a b. (a, b) -> b
snd ((Name, [BangType]) -> [BangType])
-> (Con -> (Name, [BangType])) -> Con -> [BangType]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> (Name, [BangType])
getNameAndBangTypesFromCon

-- For a type 'Tree', generates the following signature
-- transformTree ::
--     ('Data.Packed.Reader.PackedReader' '[a] r ('Data.Packed.Needs.NeedsBuilder' '[a] '[Tree a] '[] '[Tree a])) ->
--
--     ('Data.Packed.Reader.PackedReader' '[Tree a, Tree a] r ('Data.Packed.Needs.NeedsBuilder' '[Tree a, Tree a] '[Tree a] '[] '[Tree a])) ->
--     'Data.Packed.PackedReader' '[Tree a] r ('Data.Packed.Needs' '[] '[Tree a])
genTransformSignature :: [PackingFlag] -> Name -> Q Dec
genTransformSignature :: [PackingFlag] -> Name -> Q Dec
genTransformSignature [PackingFlag]
flags Name
tyName = do
    (sourceType, _) <- Name -> Q (Kind, [Name])
resolveAppliedType Name
tyName
    (TyConI (DataD _ _ _ _ cs _)) <- reify tyName
    rVar <- newName "r"
    let
        rType = Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT Name
rVar
        lambdaTypes = (\Con
c -> Con -> Kind -> Q Kind -> Maybe (Q Kind)
buildLambdaType Con
c Kind
sourceType Q Kind
rType) (Con -> Maybe (Q Kind)) -> [Con] -> [Maybe (Q Kind)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Con]
cs
        outType =
            [t|
                PackedReader
                    '[$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
sourceType)]
                    $Q Kind
rType
                    (N.Needs '[] '[$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
sourceType)])
                |]
    signature <- foldr (\Q Kind
lambda Q Kind
out -> [t|$Q Kind
lambda -> $Q Kind
out|]) outType (catMaybes lambdaTypes)
    return $ SigD (transformFName tyName) signature
  where
    -- From a constructor (say Leaf a), build type PackedTransformer a r
    buildLambdaType :: Con -> Kind -> Q Kind -> Maybe (Q Kind)
buildLambdaType Con
con Kind
ty Q Kind
restType =
        if [(Kind, Int, Bool)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Kind, Int, Bool)]
fieldType
            then Maybe (Q Kind)
forall a. Maybe a
Nothing
            else Q Kind -> Maybe (Q Kind)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Q Kind -> Maybe (Q Kind)) -> Q Kind -> Maybe (Q Kind)
forall a b. (a -> b) -> a -> b
$ do
                packedContentType <-
                    ((Kind, Int, Bool) -> Q Kind -> Q Kind)
-> Q Kind -> [(Kind, Int, Bool)] -> Q Kind
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
                        ( \(Kind
fieldTy, Int
_, Bool
needsFS) Q Kind
tys ->
                            if Bool
needsFS
                                then [t|FieldSize ': $(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
fieldTy) ': $Q Kind
tys|]
                                else [t|$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
fieldTy) ': $Q Kind
tys|]
                        )
                        [t|'[]|]
                        [(Kind, Int, Bool)]
fieldType
                [t|
                    PackedReader
                        $(return packedContentType)
                        $restType
                        (N.NeedsBuilder $(return packedContentType) '[$(return ty)] '[] '[$(return ty)])
                    |]
      where
        fieldType :: [(Kind, Int, Bool)]
fieldType = Con -> [PackingFlag] -> [(Kind, Int, Bool)]
getConFieldsIdxAndNeedsFS Con
con [PackingFlag]
flags