{-# LANGUAGE BangPatterns #-}

{-# LANGUAGE CPP #-}

{-# LANGUAGE DataKinds #-}

{-# LANGUAGE TemplateHaskell #-}

{-# LANGUAGE TypeApplications #-}

{-# LANGUAGE ViewPatterns #-}

{-# LANGUAGE DeriveDataTypeable #-}

{-# LANGUAGE RankNTypes #-}

{-# LANGUAGE ScopedTypeVariables #-}

-- | Miscellaneous Template Haskell utilities, added as needed by

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 Text.Read (readMaybe)

import TH.FixQ (fixQ)

-- | Get the 'Name' of a 'TyVarBndr'

tyVarBndrName :: TyVarBndr_ flag -> Name

tyVarBndrName = tvName

appsT :: Type -> [Type] -> Type

appsT x [] = x

appsT x (y:xs) = appsT (AppT x y) xs

-- | Breaks a type application like @A b c@ into [A, b, c]. In other

unAppsT :: Type -> [Type]

unAppsT = go []

where

go xs (AppT l x) = go (x : xs) l

go xs ty = ty : xs

-- | Given a list of types, produce the type of a tuple of

tupT :: [Q Type] -> Q Type

tupT 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.

(res, !_n) <- fixQ (\ ~(_res, n) -> foldM go (TupleT n, 0) ts)

pure res

where

go (acc, !k) ty = do

ty' <- ty

pure (acc `AppT` ty', k + 1)

-- | Given a list of types, produce the type of a promoted tuple of

promotedTupT :: [Q Type] -> Q Type

promotedTupT 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.

(res, !_n) <- fixQ (\ ~(_res, n) -> foldM go (PromotedTupleT n, 0) ts)

pure res

where

go (acc, !k) ty = do

ty' <- ty

pure (acc `AppT` ty', k + 1)

-- | Given a 'Type', returns a 'Just' value if it's a named type

typeToNamedCon :: Type -> Maybe (Name, [Type])

#if MIN_VERSION_template_haskell(2,11,0)

typeToNamedCon (InfixT l n r) = Just (n, [l, r])

typeToNamedCon (UInfixT l n r) = Just (n, [l, r])

typeToNamedCon (unAppsT -> (ConT n : args)) = Just (n, args)

typeToNamedCon _ = Nothing

-- | Expect the provided type to be an application of a regular type to

expectTyCon1 :: Name -> Type -> Q Type

expectTyCon1 expected (AppT (ConT n) x) | expected == n = return x

expectTyCon1 expected (AppT (PromotedT n) x) | expected == n = return x

expectTyCon1 expected x = fail $

"Expected " ++ pprint expected ++

", applied to one argument, but instead got " ++ pprint x ++ "."

-- | Expect the provided type to be an application of a regular type to

expectTyCon2 :: Name -> Type -> Q (Type, Type)

expectTyCon2 expected (AppT (AppT (ConT n) x) y) | expected == n = return (x, y)

expectTyCon2 expected (AppT (AppT (PromotedT n) x) y) | expected == n = return (x, y)

#if MIN_VERSION_template_haskell(2,11,0)

expectTyCon2 expected (InfixT x n y) | expected == n = return (x, y)

expectTyCon2 expected (UInfixT x n y) | expected == n = return (x, y)

expectTyCon2 expected x = fail $

"Expected " ++ pprint expected ++

", applied to two arguments, but instead got " ++ pprint x ++ "."

-- | Given a type, construct the expression (Proxy :: Proxy ty).

proxyE :: TypeQ -> ExpQ

proxyE ty = [| Proxy :: Proxy $(ty) |]

-- | Like the 'everywhere' generic traversal strategy, but skips over

everywhereButStrings :: Data a => (forall b. Data b => b -> b) -> a -> a

everywhereButStrings f =

(f . gmapT (everywhereButStrings f)) `extT` (id :: String -> String)

-- | Like the 'everywhereM' generic traversal strategy, but skips over

everywhereButStringsM :: forall a m. (Data a, Monad m) => GenericM m -> a -> m a

everywhereButStringsM f x = do

x' <- gmapM (everywhereButStringsM f) x

(f `extM` (return :: String -> m String)) x'

-- | Make a 'Name' with a 'NameS' or 'NameQ' flavour, from a 'Name' with

toSimpleName :: Name -> Name

toSimpleName = mkName . pprint

-- | Construct a plain name ('mkName') based on the given name. This is

dequalify :: Name -> Name

dequalify = mkName . nameBase

-- | Apply 'dequalify' to every type variable.

dequalifyTyVars :: Data a => a -> a

dequalifyTyVars = everywhere (id `extT` modifyType)

where

modifyType (VarT n) = VarT (dequalify n)

modifyType ty = ty

-- | Get the free type variables of a 'Type'.

freeVarsT :: Type -> [Name]

freeVarsT (ForallT tvs _ ty) = filter (`notElem` (map tyVarBndrName tvs)) (freeVarsT ty)

freeVarsT (VarT n) = [n]

freeVarsT ty = concat $ gmapQ (const [] `extQ` freeVarsT) ty

-- | Utility to conveniently handle change to 'InstanceD' API in

plainInstanceD :: Cxt -> Type -> [Dec] -> Dec

plainInstanceD =

#if MIN_VERSION_template_haskell(2,11,0)

InstanceD Nothing

InstanceD

-- | Utility to conveniently handle change to 'InstanceD' API in

fromPlainInstanceD :: Dec -> Maybe (Cxt, Type, [Dec])

#if MIN_VERSION_template_haskell(2,11,0)

fromPlainInstanceD (InstanceD _ a b c) = Just (a, b, c)

fromPlainInstanceD (InstanceD a b c) = Just (a, b, c)

fromPlainInstanceD _ = Nothing

-- | Utility to convert "Data.Typeable" 'TypeRep' to a 'Type'. Note that

typeRepToType :: TypeRep -> Q Type

typeRepToType tr = do

let (con, args) = splitTyConApp tr

modName = tyConModule con

name pn mn cn = Name (OccName cn) (NameG TcClsName (PkgName pn) (ModName mn))

conName = tyConName con

t | modName == tupleMod = TupleT $ length args

| modName == listMod && conName == listCon = ListT

| modName == typeLitsMod =

case tyConName con of

s@('"':_) -> LitT . StrTyLit $ read s

#if MIN_VERSION_template_haskell(2,18,0)

['\'', c, '\''] -> LitT $ CharTyLit c

s -> case readMaybe s of

Just n -> LitT $ NumTyLit n

_ -> error $ "Unrecognized type literal name: " ++ s

| otherwise = ConT $ name (tyConPackage con) modName conName

resultArgs <- mapM typeRepToType args

return (appsT t resultArgs)

where

typeLitsMod = tyConModule . typeRepTyCon . typeRep $ Proxy @42

tupleMod = tyConModule . typeRepTyCon . typeRep $ Proxy @((), ())

listMod = tyConModule . typeRepTyCon . typeRep $ Proxy @[()]

listCon = tyConName . typeRepTyCon . typeRep $ Proxy @[()]

-- | Hack to enable putting expressions inside 'lift'-ed TH data. For

data ExpLifter = ExpLifter

#if __GLASGOW_HASKELL__ >= 811

(forall m. Quote m => m Exp)

ExpQ

deriving (Typeable)

instance Lift ExpLifter where

lift (ExpLifter e) = e

#if MIN_VERSION_template_haskell(2,16,0)

liftTyped = error $ concat

[ "'liftTyped' is not implemented for 'ExpLifter', "

, "because it would require the generated code to have type 'ExpLifter'"

]

-- | Print splices generated by a TH splice (the printing will happen

dumpSplices :: DecsQ -> DecsQ

dumpSplices x = do

ds <- x

let code = lines (pprint ds)

reportWarning ("\n" ++ unlines (map (" " ++) code))

return ds