module Optics.TH.Internal.Utils where

import Data.Maybe
import Language.Haskell.TH
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Language.Haskell.TH.Datatype as D

import Data.Set.Optics
import Language.Haskell.TH.Optics.Internal
import Optics.Core

-- | Apply arguments to a type constructor
appsT :: TypeQ -> [TypeQ] -> TypeQ
appsT :: TypeQ -> [TypeQ] -> TypeQ
appsT = (TypeQ -> TypeQ -> TypeQ) -> TypeQ -> [TypeQ] -> TypeQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl TypeQ -> TypeQ -> TypeQ
appT

-- | Apply arguments to a function
appsE1 :: ExpQ -> [ExpQ] -> ExpQ
appsE1 :: ExpQ -> [ExpQ] -> ExpQ
appsE1 = (ExpQ -> ExpQ -> ExpQ) -> ExpQ -> [ExpQ] -> ExpQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ExpQ -> ExpQ -> ExpQ
appE

-- | Construct a tuple type given a list of types.
toTupleT :: [TypeQ] -> TypeQ
toTupleT :: [TypeQ] -> TypeQ
toTupleT [TypeQ
x] = TypeQ
x
toTupleT [TypeQ]
xs = TypeQ -> [TypeQ] -> TypeQ
appsT (Int -> TypeQ
tupleT ([TypeQ] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TypeQ]
xs)) [TypeQ]
xs

-- | Construct a tuple value given a list of expressions.
toTupleE :: [ExpQ] -> ExpQ
toTupleE :: [ExpQ] -> ExpQ
toTupleE [ExpQ
x] = ExpQ
x
toTupleE [ExpQ]
xs = [ExpQ] -> ExpQ
tupE [ExpQ]
xs

-- | Construct a tuple pattern given a list of patterns.
toTupleP :: [PatQ] -> PatQ
toTupleP :: [PatQ] -> PatQ
toTupleP [PatQ
x] = PatQ
x
toTupleP [PatQ]
xs = [PatQ] -> PatQ
tupP [PatQ]
xs

-- | Apply arguments to a type constructor.
conAppsT :: Name -> [Type] -> Type
conAppsT :: Name -> [Type] -> Type
conAppsT Name
conName = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT (Name -> Type
ConT Name
conName)

-- | Return 'Name' contained in a 'TyVarBndr'.
bndrName :: TyVarBndr -> Name
bndrName :: TyVarBndr -> Name
bndrName (PlainTV  Name
n  ) = Name
n
bndrName (KindedTV Name
n Type
_) = Name
n

-- | Generate many new names from a given base name.
newNames :: String {- ^ base name -} -> Int {- ^ count -} -> Q [Name]
newNames :: String -> Int -> Q [Name]
newNames String
base Int
n = [Q Name] -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [ String -> Q Name
newName (String
baseString -> String -> String
forall a. [a] -> [a] -> [a]
++Int -> String
forall a. Show a => a -> String
show Int
i) | Int
i <- [Int
1..Int
n] ]

-- We substitute concrete types with type variables and match them with concrete
-- types in the instance context. This significantly improves type inference as
-- GHC can match the instance more easily, but costs dependence on TypeFamilies
-- and UndecidableInstances.
eqSubst :: Type -> String -> Q (Type, Pred)
eqSubst :: Type -> String -> Q (Type, Type)
eqSubst Type
ty String
n = do
  Type
placeholder <- Name -> Type
VarT (Name -> Type) -> Q Name -> TypeQ
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Q Name
newName String
n
  (Type, Type) -> Q (Type, Type)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Type
placeholder, Type -> Type -> Type
D.equalPred Type
placeholder Type
ty)

-- | Fill in kind variables using info from datatype type parameters.
addKindVars :: D.DatatypeInfo -> Type -> Type
addKindVars :: DatatypeInfo -> Type -> Type
addKindVars = Map Name Type -> Type -> Type
forall t. SubstType t => Map Name Type -> t -> t
substType (Map Name Type -> Type -> Type)
-> (DatatypeInfo -> Map Name Type) -> DatatypeInfo -> Type -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Name, Type)] -> Map Name Type
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(Name, Type)] -> Map Name Type)
-> (DatatypeInfo -> [(Name, Type)])
-> DatatypeInfo
-> Map Name Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Type -> Maybe (Name, Type)) -> [Type] -> [(Name, Type)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe Type -> Maybe (Name, Type)
var ([Type] -> [(Name, Type)])
-> (DatatypeInfo -> [Type]) -> DatatypeInfo -> [(Name, Type)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DatatypeInfo -> [Type]
D.datatypeInstTypes
  where
    var :: Type -> Maybe (Name, Type)
var t :: Type
t@(SigT (VarT Name
n) Type
k)
      | Optic' A_Traversal NoIx Type Name -> Type -> Bool
forall k (is :: IxList) s a.
Is k A_Fold =>
Optic' k is s a -> s -> Bool
has Optic' A_Traversal NoIx Type Name
forall t. HasTypeVars t => Traversal' t Name
typeVars Type
k = (Name, Type) -> Maybe (Name, Type)
forall a. a -> Maybe a
Just (Name
n, Type
t)
      | Bool
otherwise      = Maybe (Name, Type)
forall a. Maybe a
Nothing
    var Type
_              = Maybe (Name, Type)
forall a. Maybe a
Nothing

-- | Template Haskell wants type variables declared in a forall, so
-- we find all free type variables in a given type and declare them.
quantifyType :: [TyVarBndr] -> Cxt -> Type -> Type
quantifyType :: [TyVarBndr] -> [Type] -> Type -> Type
quantifyType = Set Name -> [TyVarBndr] -> [Type] -> Type -> Type
quantifyType' Set Name
forall a. Set a
S.empty

-- | This function works like 'quantifyType' except that it takes
-- a list of variables to exclude from quantification.
quantifyType' :: S.Set Name -> [TyVarBndr] -> Cxt -> Type -> Type
quantifyType' :: Set Name -> [TyVarBndr] -> [Type] -> Type -> Type
quantifyType' Set Name
exclude [TyVarBndr]
vars [Type]
cx Type
t = [TyVarBndr] -> [Type] -> Type -> Type
ForallT [TyVarBndr]
vs [Type]
cx Type
t
  where
    vs :: [TyVarBndr]
vs = (TyVarBndr -> Bool) -> [TyVarBndr] -> [TyVarBndr]
forall a. (a -> Bool) -> [a] -> [a]
filter (\TyVarBndr
v -> TyVarBndr -> Name
bndrName TyVarBndr
v Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.notMember` Set Name
exclude)
       ([TyVarBndr] -> [TyVarBndr])
-> ([Type] -> [TyVarBndr]) -> [Type] -> [TyVarBndr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> [TyVarBndr]
D.freeVariablesWellScoped
       ([Type] -> [TyVarBndr]) -> [Type] -> [TyVarBndr]
forall a b. (a -> b) -> a -> b
$ (TyVarBndr -> Type) -> [TyVarBndr] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr -> Type
bndrToType [TyVarBndr]
vars [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ Set Type -> [Type]
forall a. Set a -> [a]
S.toList (Optic' A_Fold NoIx Type Type -> Type -> Set Type
forall k a (is :: IxList) s.
(Is k A_Fold, Ord a) =>
Optic' k is s a -> s -> Set a
setOf Optic' A_Fold NoIx Type Type
typeVarsKinded Type
t)

    bndrToType :: TyVarBndr -> Type
bndrToType (PlainTV Name
n)    = Name -> Type
VarT Name
n
    bndrToType (KindedTV Name
n Type
k) = Type -> Type -> Type
SigT (Name -> Type
VarT Name
n) Type
k

------------------------------------------------------------------------
-- Support for generating inline pragmas
------------------------------------------------------------------------

inlinePragma :: Name -> [DecQ]
inlinePragma :: Name -> [DecQ]
inlinePragma Name
methodName = [Name -> Inline -> RuleMatch -> Phases -> DecQ
pragInlD Name
methodName Inline
Inline RuleMatch
FunLike Phases
AllPhases]