{-# LANGUAGE CPP #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Miscellaneous Template Haskell utilities, added as needed by
-- packages in the th-utilities repo and elsewhere.
module TH.Utilities where

import Data.Data
import Data.Generics
import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Datatype.TyVarBndr (TyVarBndr_, tvName)

-- | Get the 'Name' of a 'TyVarBndr'
tyVarBndrName :: TyVarBndr_ flag -> Name
tyVarBndrName :: TyVarBndr_ flag -> Name
tyVarBndrName = TyVarBndr_ flag -> Name
forall flag. TyVarBndr_ flag -> Name
tvName

appsT :: Type -> [Type] -> Type
appsT :: Type -> [Type] -> Type
appsT Type
x [] = Type
x
appsT Type
x (Type
y:[Type]
xs) = Type -> [Type] -> Type
appsT (Type -> Type -> Type
AppT Type
x Type
y) [Type]
xs

-- | Breaks a type application like @A b c@ into [A, b, c]. In other
-- words, it descends leftwards down 'AppT' constructors, and yields a
-- list of the results.
unAppsT :: Type -> [Type]
unAppsT :: Type -> [Type]
unAppsT = [Type] -> Type -> [Type]
go []
  where
    go :: [Type] -> Type -> [Type]
go [Type]
xs (AppT Type
l Type
x) = [Type] -> Type -> [Type]
go (Type
x Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
xs) Type
l
    go [Type]
xs Type
ty = Type
ty Type -> [Type] -> [Type]
forall a. a -> [a] -> [a]
: [Type]
xs

-- | Given a 'Type', returns a 'Just' value if it's a named type
-- constructor applied to arguments. This value contains the name of the
-- type and a list of arguments.
typeToNamedCon :: Type -> Maybe (Name, [Type])
#if MIN_VERSION_template_haskell(2,11,0)
typeToNamedCon :: Type -> Maybe (Name, [Type])
typeToNamedCon (InfixT Type
l Name
n Type
r) = (Name, [Type]) -> Maybe (Name, [Type])
forall a. a -> Maybe a
Just (Name
n, [Type
l, Type
r])
typeToNamedCon (UInfixT Type
l Name
n Type
r) = (Name, [Type]) -> Maybe (Name, [Type])
forall a. a -> Maybe a
Just (Name
n, [Type
l, Type
r])
#endif
typeToNamedCon (Type -> [Type]
unAppsT -> (ConT Name
n : [Type]
args)) = (Name, [Type]) -> Maybe (Name, [Type])
forall a. a -> Maybe a
Just (Name
n, [Type]
args)
typeToNamedCon Type
_ = Maybe (Name, [Type])
forall a. Maybe a
Nothing

-- | Expect the provided type to be an application of a regular type to
-- one argument, otherwise fail with a message. This will also work if
-- the name is a promoted data constructor ('PromotedT').
expectTyCon1 :: Name -> Type -> Q Type
expectTyCon1 :: Name -> Type -> Q Type
expectTyCon1 Name
expected (AppT (ConT Name
n) Type
x) | Name
expected Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n = Type -> Q Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
x
expectTyCon1 Name
expected (AppT (PromotedT Name
n) Type
x) | Name
expected Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n = Type -> Q Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
x
expectTyCon1 Name
expected Type
x = String -> Q Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Type) -> String -> Q Type
forall a b. (a -> b) -> a -> b
$
    String
"Expected " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Ppr a => a -> String
pprint Name
expected String -> String -> String
forall a. [a] -> [a] -> [a]
++
    String
", applied to one argument, but instead got " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Ppr a => a -> String
pprint Type
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"."

-- | Expect the provided type to be an application of a regular type to
-- two arguments, otherwise fail with a message. This will also work if
-- the name is a promoted data constructor ('PromotedT').
expectTyCon2 :: Name -> Type -> Q (Type, Type)
expectTyCon2 :: Name -> Type -> Q (Type, Type)
expectTyCon2 Name
expected (AppT (AppT (ConT Name
n) Type
x) Type
y) | Name
expected Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n = (Type, Type) -> Q (Type, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Type
x, Type
y)
expectTyCon2 Name
expected (AppT (AppT (PromotedT Name
n) Type
x) Type
y) | Name
expected Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n = (Type, Type) -> Q (Type, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Type
x, Type
y)
#if MIN_VERSION_template_haskell(2,11,0)
expectTyCon2 Name
expected (InfixT Type
x Name
n Type
y) | Name
expected Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n = (Type, Type) -> Q (Type, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Type
x, Type
y)
expectTyCon2 Name
expected (UInfixT Type
x Name
n Type
y) | Name
expected Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
n = (Type, Type) -> Q (Type, Type)
forall (m :: * -> *) a. Monad m => a -> m a
return (Type
x, Type
y)
#endif
expectTyCon2 Name
expected Type
x = String -> Q (Type, Type)
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q (Type, Type)) -> String -> Q (Type, Type)
forall a b. (a -> b) -> a -> b
$
    String
"Expected " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Ppr a => a -> String
pprint Name
expected String -> String -> String
forall a. [a] -> [a] -> [a]
++
    String
", applied to two arguments, but instead got " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Ppr a => a -> String
pprint Type
x String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"."

-- | Given a type, construct the expression (Proxy :: Proxy ty).
proxyE :: TypeQ -> ExpQ
proxyE :: Q Type -> ExpQ
proxyE Q Type
ty = [| Proxy :: Proxy $(ty) |]

-- | Like the 'everywhere' generic traversal strategy, but skips over
-- strings. This can aid performance of TH traversals quite a bit.
everywhereButStrings :: Data a => (forall b. Data b => b -> b) -> a -> a
everywhereButStrings :: (forall b. Data b => b -> b) -> a -> a
everywhereButStrings forall b. Data b => b -> b
f =
    (a -> a
forall b. Data b => b -> b
f (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall b. Data b => b -> b) -> a -> a
forall a. Data a => (forall b. Data b => b -> b) -> a -> a
gmapT ((forall b. Data b => b -> b) -> b -> b
forall a. Data a => (forall b. Data b => b -> b) -> a -> a
everywhereButStrings forall b. Data b => b -> b
f)) (a -> a) -> (String -> String) -> a -> a
forall a b.
(Typeable a, Typeable b) =>
(a -> a) -> (b -> b) -> a -> a
`extT` (String -> String
forall a. a -> a
id :: String -> String)

-- | Like the 'everywhereM' generic traversal strategy, but skips over
-- strings. This can aid performance of TH traversals quite a bit.
everywhereButStringsM :: forall a m. (Data a, Monad m) => GenericM m -> a -> m a
everywhereButStringsM :: GenericM m -> a -> m a
everywhereButStringsM GenericM m
f a
x = do
    a
x' <- GenericM m -> a -> m a
forall a (m :: * -> *).
(Data a, Monad m) =>
(forall d. Data d => d -> m d) -> a -> m a
gmapM (GenericM m -> d -> m d
forall a (m :: * -> *). (Data a, Monad m) => GenericM m -> a -> m a
everywhereButStringsM GenericM m
f) a
x
    (a -> m a
GenericM m
f (a -> m a) -> (String -> m String) -> a -> m a
forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(a -> m a) -> (b -> m b) -> a -> m a
`extM` (String -> m String
forall (m :: * -> *) a. Monad m => a -> m a
return :: String -> m String)) a
x'

-- | Make a 'Name' with a 'NameS' or 'NameQ' flavour, from a 'Name' with
-- any 'NameFlavour'. This may change the meaning of names.
toSimpleName :: Name -> Name
toSimpleName :: Name -> Name
toSimpleName = String -> Name
mkName (String -> Name) -> (Name -> String) -> Name -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
forall a. Ppr a => a -> String
pprint

-- | Construct a plain name ('mkName') based on the given name. This is
-- useful for cases where TH doesn't expect a unique name.
dequalify :: Name -> Name
dequalify :: Name -> Name
dequalify = String -> Name
mkName (String -> Name) -> (Name -> String) -> Name -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase

-- | Apply 'dequalify' to every type variable.
dequalifyTyVars :: Data a => a -> a
dequalifyTyVars :: a -> a
dequalifyTyVars = (forall b. Data b => b -> b) -> forall b. Data b => b -> b
everywhere (a -> a
forall a. a -> a
id (a -> a) -> (Type -> Type) -> a -> a
forall a b.
(Typeable a, Typeable b) =>
(a -> a) -> (b -> b) -> a -> a
`extT` Type -> Type
modifyType)
  where
    modifyType :: Type -> Type
modifyType (VarT Name
n) = Name -> Type
VarT (Name -> Name
dequalify Name
n)
    modifyType Type
ty = Type
ty

-- | Get the free type variables of a 'Type'.
freeVarsT :: Type -> [Name]
freeVarsT :: Type -> [Name]
freeVarsT (ForallT [TyVarBndr_ flag]
tvs [Type]
_ Type
ty) = (Name -> Bool) -> [Name] -> [Name]
forall a. (a -> Bool) -> [a] -> [a]
filter (Name -> [Name] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` ((TyVarBndr_ flag -> Name) -> [TyVarBndr_ flag] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr_ flag -> Name
forall flag. TyVarBndr_ flag -> Name
tyVarBndrName [TyVarBndr_ flag]
tvs)) (Type -> [Name]
freeVarsT Type
ty)
freeVarsT (VarT Name
n) = [Name
n]
freeVarsT Type
ty = [[Name]] -> [Name]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Name]] -> [Name]) -> [[Name]] -> [Name]
forall a b. (a -> b) -> a -> b
$ (forall d. Data d => d -> [Name]) -> Type -> [[Name]]
forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
gmapQ ([Name] -> d -> [Name]
forall a b. a -> b -> a
const [] (d -> [Name]) -> (Type -> [Name]) -> d -> [Name]
forall a b q.
(Typeable a, Typeable b) =>
(a -> q) -> (b -> q) -> a -> q
`extQ` Type -> [Name]
freeVarsT) Type
ty

-- | Utility to conveniently handle change to 'InstanceD' API in
-- template-haskell-2.11.0
plainInstanceD :: Cxt -> Type -> [Dec] -> Dec
plainInstanceD :: [Type] -> Type -> [Dec] -> Dec
plainInstanceD =
#if MIN_VERSION_template_haskell(2,11,0)
    Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
forall a. Maybe a
Nothing
#else
    InstanceD
#endif

-- | Utility to conveniently handle change to 'InstanceD' API in
-- template-haskell-2.11.0
fromPlainInstanceD :: Dec -> Maybe (Cxt, Type, [Dec])
#if MIN_VERSION_template_haskell(2,11,0)
fromPlainInstanceD :: Dec -> Maybe ([Type], Type, [Dec])
fromPlainInstanceD (InstanceD Maybe Overlap
_ [Type]
a Type
b [Dec]
c) = ([Type], Type, [Dec]) -> Maybe ([Type], Type, [Dec])
forall a. a -> Maybe a
Just ([Type]
a, Type
b, [Dec]
c)
#else
fromPlainInstanceD (InstanceD a b c) = Just (a, b, c)
#endif
fromPlainInstanceD Dec
_ = Maybe ([Type], Type, [Dec])
forall a. Maybe a
Nothing

-- | Utility to convert "Data.Typeable" 'TypeRep' to a 'Type'. Note that
-- this function is known to not yet work for many cases, but it does
-- work for normal user datatypes. In future versions this function
-- might have better behavior.
typeRepToType :: TypeRep -> Q Type
typeRepToType :: TypeRep -> Q Type
typeRepToType TypeRep
tr = do
    let (TyCon
con, [TypeRep]
args) = TypeRep -> (TyCon, [TypeRep])
splitTyConApp TypeRep
tr
        name :: Name
name = OccName -> NameFlavour -> Name
Name (String -> OccName
OccName (TyCon -> String
tyConName TyCon
con)) (NameSpace -> PkgName -> ModName -> NameFlavour
NameG NameSpace
TcClsName (String -> PkgName
PkgName (TyCon -> String
tyConPackage TyCon
con)) (String -> ModName
ModName (TyCon -> String
tyConModule TyCon
con)))
    [Type]
resultArgs <- (TypeRep -> Q Type) -> [TypeRep] -> Q [Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeRep -> Q Type
typeRepToType [TypeRep]
args
    Type -> Q Type
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> [Type] -> Type
appsT (Name -> Type
ConT Name
name) [Type]
resultArgs)

-- | Hack to enable putting expressions inside 'lift'-ed TH data. For
-- example, you could do
--
-- @
--     main = print $(lift [ExpLifter [e| 1 + 1 |],  ExpLifter [e| 2 |]])
-- @
--
-- Here, 'lift' is working on a value of type @[ExpLifter]@. The code
-- generated by 'lift' constructs a list with the 'ExpLifter'
-- expressions providing the element values.
--
-- Without 'ExpLifter', 'lift' tends to just generate code involving
-- data construction. With 'ExpLifter', you can put more complicated
-- expression into this construction.
--
-- Note that this cannot be used in typed quotes, because 'liftTyped'
-- will throw an exception. This is because this hack is incompatible
-- with the type of 'liftTyped', as it would require the generated
-- code to have type 'ExpLifter'.
data ExpLifter = ExpLifter
#if __GLASGOW_HASKELL__ >= 811
  (forall m. Quote m => m Exp)
#else
  ExpQ
#endif
  deriving (Typeable)

instance Lift ExpLifter where
  lift :: ExpLifter -> ExpQ
lift (ExpLifter ExpQ
e) = ExpQ
e
#if MIN_VERSION_template_haskell(2,16,0)
  liftTyped :: ExpLifter -> Q (TExp ExpLifter)
liftTyped = String -> ExpLifter -> Q (TExp ExpLifter)
forall a. HasCallStack => String -> a
error (String -> ExpLifter -> Q (TExp ExpLifter))
-> String -> ExpLifter -> Q (TExp ExpLifter)
forall a b. (a -> b) -> a -> b
$ [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    [ String
"'liftTyped' is not implemented for 'ExpLifter', "
    , String
"because it would require the generated code to have type 'ExpLifter'"
    ]
#endif

-- | Print splices generated by a TH splice (the printing will happen
-- during compilation, as a GHC warning). Useful for debugging.
--
-- For instance, you can dump splices generated with 'makeLenses' by
-- replacing a top-level invocation of 'makeLenses' in your code with:
--
-- @dumpSplices $ makeLenses ''Foo@
dumpSplices :: DecsQ -> DecsQ
dumpSplices :: DecsQ -> DecsQ
dumpSplices DecsQ
x = do
  [Dec]
ds <- DecsQ
x
  let code :: [String]
code = String -> [String]
lines ([Dec] -> String
forall a. Ppr a => a -> String
pprint [Dec]
ds)
  String -> Q ()
reportWarning (String
"\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ [String] -> String
unlines ((String -> String) -> [String] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (String
"    " String -> String -> String
forall a. [a] -> [a] -> [a]
++) [String]
code))
  [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return [Dec]
ds