1module Optics.TH.Internal.Utils where
2
3import Data.Maybe
4import Language.Haskell.TH
5import qualified Data.Map as M
6import qualified Data.Set as S
7import qualified Language.Haskell.TH.Datatype as D
8
9import Data.Set.Optics
10import Language.Haskell.TH.Optics.Internal
11import Optics.Core
12
13-- | Apply arguments to a type constructor
14appsT :: TypeQ -> [TypeQ] -> TypeQ
15appsT = foldl appT
16
17-- | Apply arguments to a function
18appsE1 :: ExpQ -> [ExpQ] -> ExpQ
19appsE1 = foldl appE
20
21-- | Construct a tuple type given a list of types.
22toTupleT :: [TypeQ] -> TypeQ
23toTupleT [x] = x
24toTupleT xs = appsT (tupleT (length xs)) xs
25
26-- | Construct a tuple value given a list of expressions.
27toTupleE :: [ExpQ] -> ExpQ
28toTupleE [x] = x
29toTupleE xs = tupE xs
30
31-- | Construct a tuple pattern given a list of patterns.
32toTupleP :: [PatQ] -> PatQ
33toTupleP [x] = x
34toTupleP xs = tupP xs
35
36-- | Apply arguments to a type constructor.
37conAppsT :: Name -> [Type] -> Type
38conAppsT conName = foldl AppT (ConT conName)
39
40-- | Return 'Name' contained in a 'TyVarBndr'.
41bndrName :: TyVarBndr -> Name
42bndrName (PlainTV  n  ) = n
43bndrName (KindedTV n _) = n
44
45-- | Generate many new names from a given base name.
46newNames :: String {- ^ base name -} -> Int {- ^ count -} -> Q [Name]
47newNames base n = sequence [ newName (base++show i) | i <- [1..n] ]
48
49-- We substitute concrete types with type variables and match them with concrete
50-- types in the instance context. This significantly improves type inference as
51-- GHC can match the instance more easily, but costs dependence on TypeFamilies
52-- and UndecidableInstances.
53eqSubst :: Type -> String -> Q (Type, Pred)
54eqSubst ty n = do
55  placeholder <- VarT <$> newName n
56  pure (placeholder, D.equalPred placeholder ty)
57
58-- | Fill in kind variables using info from datatype type parameters.
59addKindVars :: D.DatatypeInfo -> Type -> Type
60addKindVars = substType . M.fromList . mapMaybe var . D.datatypeInstTypes
61  where
62    var t@(SigT (VarT n) k)
63      | has typeVars k = Just (n, t)
64      | otherwise      = Nothing
65    var _              = Nothing
66
67-- | Template Haskell wants type variables declared in a forall, so
68-- we find all free type variables in a given type and declare them.
69quantifyType :: [TyVarBndr] -> Cxt -> Type -> Type
70quantifyType = quantifyType' S.empty
71
72-- | This function works like 'quantifyType' except that it takes
73-- a list of variables to exclude from quantification.
74quantifyType' :: S.Set Name -> [TyVarBndr] -> Cxt -> Type -> Type
75quantifyType' exclude vars cx t = ForallT vs cx t
76  where
77    vs = filter (\v -> bndrName v `S.notMember` exclude)
78       . D.freeVariablesWellScoped
79       $ map bndrToType vars ++ S.toList (setOf typeVarsKinded t)
80
81    bndrToType (PlainTV n)    = VarT n
82    bndrToType (KindedTV n k) = SigT (VarT n) k
83
84------------------------------------------------------------------------
85-- Support for generating inline pragmas
86------------------------------------------------------------------------
87
88inlinePragma :: Name -> [DecQ]
89inlinePragma methodName = [pragInlD methodName Inline FunLike AllPhases]
90