{-# LANGUAGE BangPatterns #-}
{-# 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 Control.Monad (foldM)
import Data.Data
import Data.Generics
import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Datatype.TyVarBndr (TyVarBndr_, tvName)
import TH.FixQ (fixQ)

-- | Get the 'Name' of a 'TyVarBndr'
tyVarBndrName :: TyVarBndr_ flag -> Name
tyVarBndrName :: forall flag. TyVarBndr_ flag -> Name
tyVarBndrName = 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 forall a. a -> [a] -> [a]
: [Type]
xs) Type
l
    go [Type]
xs Type
ty = Type
ty forall a. a -> [a] -> [a]
: [Type]
xs

-- | Given a list of types, produce the type of a tuple of
-- those types. This is analogous to 'tupE' and 'tupP'.
--
-- @
-- tupT [[t|Int|], [t|Char|], [t|Bool]] = [t| (Int, Char, Bool) |]
-- @
--
-- @since FIXME
tupT :: [Q Type] -> Q Type
tupT :: [Q Type] -> Q Type
tupT [Q Type]
ts = do
  -- We build the expression with a thunk inside that will be filled in with
  -- the length of the list once that's been determined. This works
  -- efficiently (in one pass) because TH.Type is rather lazy.
  (Type
res, !Int
_n) <- forall a. (a -> Q a) -> Q a
fixQ (\ ~(Type
_res, Int
n) -> forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {m :: * -> *} {b}.
(Monad m, Num b) =>
(Type, b) -> m Type -> m (Type, b)
go (Int -> Type
TupleT Int
n, Int
0) [Q Type]
ts)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
res
  where
    go :: (Type, b) -> m Type -> m (Type, b)
go (Type
acc, !b
k) m Type
ty = do
      Type
ty' <- m Type
ty
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
acc Type -> Type -> Type
`AppT` Type
ty', b
k forall a. Num a => a -> a -> a
+ b
1)

-- | Given a list of types, produce the type of a promoted tuple of
-- those types. This is analogous to 'tupE' and 'tupP'.
--
-- @
-- promotedTupT [[t|3|], [t| 'True|], [t|Bool]] = [t| '(3, 'True, Bool) |]
-- @
--
-- @since FIXME
promotedTupT :: [Q Type] -> Q Type
promotedTupT :: [Q Type] -> Q Type
promotedTupT [Q Type]
ts = do
  -- We build the expression with a thunk inside that will be filled in with
  -- the length of the list once that's been determined. This works
  -- efficiently (in one pass) because TH.Type is rather lazy.
  (Type
res, !Int
_n) <- forall a. (a -> Q a) -> Q a
fixQ (\ ~(Type
_res, Int
n) -> forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM forall {m :: * -> *} {b}.
(Monad m, Num b) =>
(Type, b) -> m Type -> m (Type, b)
go (Int -> Type
PromotedTupleT Int
n, Int
0) [Q Type]
ts)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure Type
res
  where
    go :: (Type, b) -> m Type -> m (Type, b)
go (Type
acc, !b
k) m Type
ty = do
      Type
ty' <- m Type
ty
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
acc Type -> Type -> Type
`AppT` Type
ty', b
k forall a. Num a => a -> a -> a
+ b
1)

-- | 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) = forall a. a -> Maybe a
Just (Name
n, [Type
l, Type
r])
typeToNamedCon (UInfixT Type
l Name
n Type
r) = forall a. a -> Maybe a
Just (Name
n, [Type
l, Type
r])
#endif
typeToNamedCon (Type -> [Type]
unAppsT -> (ConT Name
n : [Type]
args)) = forall a. a -> Maybe a
Just (Name
n, [Type]
args)
typeToNamedCon 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 forall a. Eq a => a -> a -> Bool
== Name
n = forall (m :: * -> *) a. Monad m => a -> m a
return Type
x
expectTyCon1 Name
expected (AppT (PromotedT Name
n) Type
x) | Name
expected forall a. Eq a => a -> a -> Bool
== Name
n = forall (m :: * -> *) a. Monad m => a -> m a
return Type
x
expectTyCon1 Name
expected Type
x = forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$
    String
"Expected " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> String
pprint Name
expected forall a. [a] -> [a] -> [a]
++
    String
", applied to one argument, but instead got " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> String
pprint Type
x 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 forall a. Eq a => a -> a -> Bool
== Name
n = 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 forall a. Eq a => a -> a -> Bool
== Name
n = 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 forall a. Eq a => a -> a -> Bool
== Name
n = 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 forall a. Eq a => a -> a -> Bool
== Name
n = forall (m :: * -> *) a. Monad m => a -> m a
return (Type
x, Type
y)
#endif
expectTyCon2 Name
expected Type
x = forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$
    String
"Expected " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> String
pprint Name
expected forall a. [a] -> [a] -> [a]
++
    String
", applied to two arguments, but instead got " forall a. [a] -> [a] -> [a]
++ forall a. Ppr a => a -> String
pprint Type
x 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 a. Data a => (forall b. Data b => b -> b) -> a -> a
everywhereButStrings forall b. Data b => b -> b
f =
    (forall b. Data b => b -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Data a => (forall b. Data b => b -> b) -> a -> a
gmapT (forall a. Data a => (forall b. Data b => b -> b) -> a -> a
everywhereButStrings forall b. Data b => b -> b
f)) forall a b.
(Typeable a, Typeable b) =>
(a -> a) -> (b -> b) -> a -> a
`extT` (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 :: forall a (m :: * -> *). (Data a, Monad m) => GenericM m -> a -> m a
everywhereButStringsM GenericM m
f a
x = do
    a
x' <- forall a (m :: * -> *).
(Data a, Monad m) =>
(forall d. Data d => d -> m d) -> a -> m a
gmapM (forall a (m :: * -> *). (Data a, Monad m) => GenericM m -> a -> m a
everywhereButStringsM GenericM m
f) a
x
    (GenericM m
f forall (m :: * -> *) a b.
(Monad m, Typeable a, Typeable b) =>
(a -> m a) -> (b -> m b) -> a -> m a
`extM` (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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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 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 :: forall b. Data b => b -> b
dequalifyTyVars = (forall b. Data b => b -> b) -> forall b. Data b => b -> b
everywhere (forall a. a -> a
id 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 Specificity]
tvs [Type]
_ Type
ty) = forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` (forall a b. (a -> b) -> [a] -> [b]
map forall flag. TyVarBndr_ flag -> Name
tyVarBndrName [TyVarBndr Specificity]
tvs)) (Type -> [Name]
freeVarsT Type
ty)
freeVarsT (VarT Name
n) = [Name
n]
freeVarsT Type
ty = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ forall a u. Data a => (forall d. Data d => d -> u) -> a -> [u]
gmapQ (forall a b. a -> b -> a
const [] 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 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) = 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
_ = 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 <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TypeRep -> Q Type
typeRepToType [TypeRep]
args
    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 :: forall (m :: * -> *). Quote m => ExpLifter -> m Exp
lift (ExpLifter forall (m :: * -> *). Quote m => m Exp
e) = forall (m :: * -> *). Quote m => m Exp
e
#if MIN_VERSION_template_haskell(2,16,0)
  liftTyped :: forall (m :: * -> *). Quote m => ExpLifter -> Code m ExpLifter
liftTyped = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ 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 (forall a. Ppr a => a -> String
pprint [Dec]
ds)
  String -> Q ()
reportWarning (String
"\n" forall a. [a] -> [a] -> [a]
++ [String] -> String
unlines (forall a b. (a -> b) -> [a] -> [b]
map (String
"    " forall a. [a] -> [a] -> [a]
++) [String]
code))
  forall (m :: * -> *) a. Monad m => a -> m a
return [Dec]
ds