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 = case 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
<$> (Name, [BangType]) -> [BangType]
forall a b. (a, b) -> b
snd (Con -> (Name, [BangType])
getNameAndBangTypesFromCon Con
con) of
        [] -> Maybe (Q Kind)
forall a. Maybe a
Nothing
        Cxt
constructorTypeNames -> 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 <-
                ((Int, Kind) -> Q Kind -> Q Kind)
-> Q Kind -> [(Int, Kind)] -> 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
                    ( \(Int
i, Kind
x) Q Kind
xs ->
                        if (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
|| (PackingFlag
SkipLastFieldSize 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
&& Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1))
                            then [t|'[FieldSize, $(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
x)] :++: $Q Kind
xs|]
                            else [t|$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
x) ': $Q Kind
xs|]
                    )
                    [t|'[]|]
                    ([(Int, Kind)] -> Q Kind) -> [(Int, Kind)] -> Q Kind
forall a b. (a -> b) -> a -> b
$ [Int] -> Cxt -> [(Int, Kind)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int
0 .. Cxt -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Cxt
constructorTypeNames]) Cxt
constructorTypeNames
            [t|
                PackedReader
                    $(return packedContentType)
                    $restType
                    (N.NeedsBuilder $(return packedContentType) '[$(return ty)] '[] '[$(return ty)])
                |]