module TypeMachine.TH.Is (isClassName, deriveIs, defineIs) where

import Control.Monad (MonadPlus (mzero), forM)
import qualified Data.Map.Strict as Map
import Language.Haskell.TH hiding (Type, reifyType)
import qualified Language.Haskell.TH as TH
import Text.Printf
import TypeMachine.TH.Internal.Utils
import TypeMachine.Type (fields, reifyType)

-- | Get the name of the 'Is' class generated for the given type
--
-- @
-- > isClassName ''User
--   IsUser
-- @
isClassName :: Name -> Name
isClassName :: Name -> Name
isClassName = String -> Name
mkName (String -> Name) -> (Name -> String) -> Name -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"Is" String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> String) -> (Name -> String) -> Name -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
capitalize (String -> String) -> (Name -> String) -> Name -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase

-- | Get the name of the 'to' function generated for the given type
--
-- @
-- > toFuncName ''User
--   toUser
-- @
toFuncName :: Name -> Name
toFuncName :: Name -> Name
toFuncName = String -> Name
mkName (String -> Name) -> (Name -> String) -> Name -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"to" String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> String) -> (Name -> String) -> Name -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
capitalize (String -> String) -> (Name -> String) -> Name -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase

-- | Returns the declaration of the instance of 'Is' for a given type
--
-- @
--  > deriveIs ''Animal ''Dog
--
--    instance IsAnimal Dog where
--       ...
-- @
deriveIs :: Name -> Name -> Q [Dec]
deriveIs :: Name -> Name -> Q [Dec]
deriveIs Name
sourceTypeName Name
destTypeName = do
    destFields <- Type -> Map String BangType
fields (Type -> Map String BangType) -> Q Type -> Q (Map String BangType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> Q Type
reifyType Name
destTypeName
    sourceFields <- fields <$> reifyType sourceTypeName
    let className = String -> Name
mkName (String
"Is" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
sourceTypeName)
    classFuncs <- fmap concat $ forM (zip [0 ..] $ Map.toList sourceFields) $ \(Int
i, (String
n, (Bang
_, Type
t))) ->
        case String -> Map String BangType -> Maybe BangType
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup String
n Map String BangType
destFields of
            Just BangType
_ -> do
                getter <- String -> Q Dec
forall {m :: * -> *}. Quote m => String -> m Dec
fieldToGetter String
n
                setter <- fieldToSetter (length destFields) i n
                return [getter, setter]
            Maybe BangType
Nothing ->
                Q Bool -> Q [Dec] -> Q [Dec] -> Q [Dec]
forall {m :: * -> *} {b}. Monad m => m Bool -> m b -> m b -> m b
ifM
                    (Type -> Q Bool
fieldIsOptional Type
t)
                    ( do
                        getter <- String -> Q Dec
forall {m :: * -> *}. Quote m => String -> m Dec
fieldNameToMemptyFunDec String
n
                        setter <- fieldNameToNoopFunDec n
                        return [getter, setter]
                    )
                    ( String -> Q [Dec]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail
                        ( String -> String -> String -> String -> String -> String
forall r. PrintfType r => String -> r
printf
                            String
"Type-Machine Error: Cannot define instance of %s for %s. Field '%s' is missing in %s "
                            (Name -> String
nameBase Name
className)
                            String
destTypeStr
                            String
n
                            String
destTypeStr
                        )
                    )
    let inlinePragmas = [String] -> [Dec]
mkInlinePragmas ([String] -> [Dec]) -> [String] -> [Dec]
forall a b. (a -> b) -> a -> b
$ Map String BangType -> [String]
forall k a. Map k a -> [k]
Map.keys Map String BangType
sourceFields
        instanceDec =
            Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD
                Maybe Overlap
forall a. Maybe a
Nothing
                []
                (Type -> Type -> Type
AppT (Name -> Type
ConT Name
className) (Name -> Type
ConT Name
destTypeName))
                ([Dec]
inlinePragmas [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [Dec]
classFuncs)
    return [instanceDec]
  where
    destTypeStr :: String
destTypeStr = Name -> String
nameBase Name
destTypeName
    fieldNameToMemptyFunDec :: String -> m Dec
fieldNameToMemptyFunDec String
n =
        Name -> [m Clause] -> m Dec
forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
funD (String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String -> String
fieldNameToIsGetter String
n) [[m Pat] -> m Body -> [m Dec] -> m Clause
forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
clause [] (m Exp -> m Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|const mzero|]) []]

    fieldNameToNoopFunDec :: String -> m Dec
fieldNameToNoopFunDec String
n =
        Name -> [m Clause] -> m Dec
forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
funD
            (String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String -> String
fieldNameToIsSetter String
n)
            [[m Pat] -> m Body -> [m Dec] -> m Clause
forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
clause [m Pat
forall (m :: * -> *). Quote m => m Pat
wildP, Name -> m Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
inputObjName] (m Exp -> m Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (m Exp -> m Body) -> m Exp -> m Body
forall a b. (a -> b) -> a -> b
$ Name -> m Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
inputObjName) []]
      where
        inputObjName :: Name
inputObjName = String -> Name
mkName String
"x"
    fieldToGetter :: String -> m Dec
fieldToGetter String
n = do
        let funName :: Name
funName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String -> String
fieldNameToIsGetter String
n
            resName :: Name
resName = String -> Name
mkName String
"res"
            expr :: m Exp
expr = [|$(Name -> m Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
resName)|]
        -- Note: using destTypeName makes Q think that we use the type, not the constructor
        Name -> [m Clause] -> m Dec
forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
funD
            Name
funName
            [[m Pat] -> m Body -> [m Dec] -> m Clause
forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
clause [Pat -> m Pat
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Pat -> m Pat) -> Pat -> m Pat
forall a b. (a -> b) -> a -> b
$ Name -> [FieldPat] -> Pat
RecP (String -> Name
mkName String
destTypeStr) [(String -> Name
mkName String
n, Name -> Pat
VarP Name
resName)]] (m Exp -> m Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB m Exp
expr) []]

    fieldToSetter :: a -> a -> String -> m Dec
fieldToSetter a
fieldCount a
fieldPos String
fieldName = do
        fieldsNames <-
            [a] -> (a -> m Name) -> m [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM
                [a
0 .. (a
fieldCount a -> a -> a
forall a. Num a => a -> a -> a
- a
1)]
                ( \a
i ->
                    if a
i a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
fieldPos
                        then Name -> m Name
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> m Name) -> Name -> m Name
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"_"
                        else String -> m Name
forall (m :: * -> *). Quote m => String -> m Name
newName (String -> m Name) -> String -> m Name
forall a b. (a -> b) -> a -> b
$ String
"f" String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
i
                )
        let funName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String -> String
fieldNameToIsSetter String
fieldName
        let newValueName = String -> Name
mkName String
"new"
        let patt = Name -> Cxt -> [Pat] -> Pat
ConP (String -> Name
mkName String
destTypeStr) [] (Name -> Pat
VarP (Name -> Pat) -> [Name] -> [Pat]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
fieldsNames)
        let body =
                (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
                    ( \Exp
res Name
f ->
                        Exp
res Exp -> Exp -> Exp
`AppE` case Name -> String
nameBase Name
f of
                            String
"_" -> Name -> Exp
VarE Name
newValueName
                            String
_ -> Name -> Exp
VarE Name
f
                    )
                    (Name -> Exp
ConE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
destTypeStr)
                    [Name]
fieldsNames

        funD funName [clause [varP newValueName, return $ patt] (normalB $ return body) []]
    -- TODO handle non-parametric monadplus-es
    -- Returns true is field is instance of Monad plus
    fieldIsOptional :: TH.Type -> Q Bool
    fieldIsOptional :: Type -> Q Bool
fieldIsOptional (AppT Type
t Type
_) = Name -> Cxt -> Q Bool
isInstance ''MonadPlus [Type
t]
    fieldIsOptional Type
_ = Bool -> Q Bool
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
    ifM :: m Bool -> m b -> m b -> m b
ifM m Bool
mbool m b
t m b
f = do bool <- m Bool
mbool; if bool then t else f
    mkInlinePragmas :: [String] -> [Dec]
    mkInlinePragmas :: [String] -> [Dec]
mkInlinePragmas [String]
fieldNames = Name -> Dec
mkInlinePragma (Name -> Dec) -> (String -> Name) -> String -> Dec
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Name
mkName (String -> Dec) -> [String] -> [Dec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
setterAndGetterNames
      where
        setterAndGetterNames :: [String]
setterAndGetterNames =
            (String -> [String] -> [String])
-> [String] -> [String] -> [String]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
                (\String
field [String]
rest -> String -> String
fieldNameToIsSetter String
field String -> [String] -> [String]
forall a. a -> [a] -> [a]
: String -> String
fieldNameToIsGetter String
field String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
rest)
                []
                [String]
fieldNames
        mkInlinePragma :: Name -> Dec
mkInlinePragma Name
fName = Pragma -> Dec
PragmaD (Pragma -> Dec) -> Pragma -> Dec
forall a b. (a -> b) -> a -> b
$ Name -> Inline -> RuleMatch -> Phases -> Pragma
InlineP Name
fName Inline
Inline RuleMatch
FunLike Phases
AllPhases

----- Definition

-- | Define the 'Is' class for the given type and generate the 'To' function
--
-- @
--  > data User = User { id :: Int, name :: String }
--  > defineIs ''User
--
--  class IsUser a where
--       getId :: a -> Int
--       getName :: a -> String
--
--       setId :: Int -> a -> a
--       setName :: String -> a -> a
--
--       toUser :: (IsUser a) => a -> User
--       toUser a = User (getId a) (getName a)
--
--  instance IsUser User where
--       ...
--
-- @
defineIs :: Name -> Q [Dec]
defineIs :: Name -> Q [Dec]
defineIs Name
tyName = do
    ty <- Name -> Q Type
reifyType Name
tyName
    classTypeVar <- newName "a"
    getters <- mapM (vbtToGetter classTypeVar) (Map.toList $ fields ty)
    setters <- mapM (vbtToSetter classTypeVar) (Map.toList $ fields ty)
    to <- defineTo tyName classTypeVar
    isItself <- deriveIs tyName tyName
    return $
        ClassD
            []
            (isClassName tyName)
            [PlainTV classTypeVar BndrReq]
            []
            (getters ++ setters ++ to)
            : isItself
  where
    -- vbtToGetter a id Int == getId :: a -> Int
    vbtToGetter :: Name -> (String, BangType) -> Q Dec
    vbtToGetter :: Name -> (String, BangType) -> Q Dec
vbtToGetter Name
classtypeVar (String
n, (Bang
_, Type
t)) =
        let
            memberName :: Name
memberName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String -> String
fieldNameToIsGetter String
n
         in
            Name -> Q Type -> Q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
memberName [t|$(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT Name
classtypeVar) -> $(Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t)|]

    -- vbtToSetter a id Int == setId ::  Int -> a -> a
    vbtToSetter :: Name -> (String, BangType) -> Q Dec
    vbtToSetter :: Name -> (String, BangType) -> Q Dec
vbtToSetter Name
classtypeVar (String
n, (Bang
_, Type
t)) =
        let
            memberName :: Name
memberName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String -> String
fieldNameToIsSetter String
n
         in
            Name -> Q Type -> Q Dec
forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
memberName [t|$(Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t) -> $(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT Name
classtypeVar) -> $(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT Name
classtypeVar)|]

-- | Generate the 'To' function
--
-- @
--  > data User = User { id :: Int, name :: String }
--
--  toUser :: (IsUser a) => a -> User
--  toUser from = User (getId from) (getName from)
-- @
defineTo :: Name -> Name -> Q [Dec]
defineTo :: Name -> Name -> Q [Dec]
defineTo Name
tyName Name
tyVarName = do
    ty <- Name -> Q Type
reifyType Name
tyName
    toFuncType <-
        sigD
            toName
            [t|$(varT tyVarName) -> $(conT tyName)|]
    toFuncBody <-
        let
            from = String -> Name
mkName String
"from"
            app =
                (Q Exp -> String -> Q Exp) -> Q Exp -> [String] -> 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
r String
n -> [|$Q Exp
r ($(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE (Name -> Q Exp) -> Name -> Q Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String -> String
fieldNameToIsGetter String
n) $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
from))|])
                    (Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE (Name -> Q Exp) -> Name -> Q Exp
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase Name
tyName)
                    (Map String BangType -> [String]
forall k a. Map k a -> [k]
Map.keys (Map String BangType -> [String])
-> Map String BangType -> [String]
forall a b. (a -> b) -> a -> b
$ Type -> Map String BangType
fields Type
ty)
         in
            funD (toFuncName tyName) [clause [varP from] (normalB app) []]
    return [toFuncType, toFuncBody]
  where
    toName :: Name
toName = Name -> Name
toFuncName Name
tyName

-- Internal
fieldNameToIsGetter :: String -> String
fieldNameToIsGetter :: String -> String
fieldNameToIsGetter = (String
"get" String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
capitalize

fieldNameToIsSetter :: String -> String
fieldNameToIsSetter :: String -> String
fieldNameToIsSetter = (String
"set" String -> String -> String
forall a. [a] -> [a] -> [a]
++) (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
capitalize