{-# LANGUAGE CPP #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
module TH.ReifySimple
(
TypeInfo, reifyType, infoToType
, reifyTypeNoDataKinds, infoToTypeNoDataKinds
, DataType(..), reifyDataType, infoToDataType
, DataCon(..), reifyDataCon, infoToDataCon, typeToDataCon
, DataFamily(..), DataInst(..), reifyDataFamily, infoToDataFamily
, TypeFamily(..), TypeInst(..), reifyTypeFamily, infoToTypeFamily
, conToDataCons
, reifyDataTypeSubstituted
) where
import Control.Applicative
import Data.Data (Data, gmapT)
import Data.Generics.Aliases (extT)
import qualified Data.Map as M
import Data.Typeable (Typeable)
import GHC.Generics (Generic)
import Language.Haskell.TH
import Language.Haskell.TH.Instances ()
import TH.Utilities
data TypeInfo
= DataTypeInfo DataType
| DataFamilyInfo DataFamily
| TypeFamilyInfo TypeFamily
| LiftedDataConInfo DataCon
reifyType :: Name -> Q TypeInfo
reifyType name = do
info <- reify name
mres <- infoToType info
case mres of
Just res -> return res
Nothing -> fail $
"Expected to reify a data type, data family, or type family. Instead got:\n" ++
pprint info
infoToType :: Info -> Q (Maybe TypeInfo)
infoToType info =
case (infoToTypeNoDataKinds info, infoToDataCon info) of
(Just result, _) -> return (Just result)
(Nothing, Just dc) -> do
#if MIN_VERSION_template_haskell(2,11,0)
dataKindsEnabled <- isExtEnabled DataKinds
#else
reportWarning $
"For " ++ pprint (dcName dc) ++
", assuming DataKinds is on, and yielding LiftedDataConInfo."
let dataKindsEnabled = True
#endif
return $ if dataKindsEnabled then Just (LiftedDataConInfo dc) else Nothing
(Nothing, Nothing) -> return Nothing
reifyTypeNoDataKinds :: Name -> Q (Maybe TypeInfo)
reifyTypeNoDataKinds = fmap infoToTypeNoDataKinds . reify
infoToTypeNoDataKinds :: Info -> Maybe TypeInfo
infoToTypeNoDataKinds info =
(DataTypeInfo <$> infoToDataType info) <|>
(DataFamilyInfo <$> infoToDataFamily info) <|>
(TypeFamilyInfo <$> infoToTypeFamily info)
data DataType = DataType
{ dtName :: Name
, dtTvs :: [Name]
, dtCxt :: Cxt
, dtCons :: [DataCon]
} deriving (Eq, Show, Ord, Data, Typeable, Generic)
data DataCon = DataCon
{ dcName :: Name
, dcTvs :: [Name]
, dcCxt :: Cxt
, dcFields :: [(Maybe Name, Type)]
} deriving (Eq, Show, Ord, Data, Typeable, Generic)
data DataFamily = DataFamily
{ dfName :: Name
, dfTvs :: [Name]
, dfInsts :: [DataInst]
} deriving (Eq, Show, Ord, Data, Typeable, Generic)
data DataInst = DataInst
{ diName :: Name
, diCxt :: Cxt
, diParams :: [Type]
, diCons :: [DataCon]
} deriving (Eq, Show, Ord, Data, Typeable, Generic)
data TypeFamily = TypeFamily
{ tfName :: Name
, tfTvs :: [Name]
, tfInsts :: [TypeInst]
} deriving (Eq, Show, Ord, Data, Typeable, Generic)
data TypeInst = TypeInst
{ tiName :: Name
, tiParams :: [Type]
, tiType :: Type
} deriving (Eq, Show, Ord, Data, Typeable, Generic)
reifyDataType :: Name -> Q DataType
reifyDataType name = do
info <- reify name
case infoToDataType info of
Nothing -> fail $ "Expected to reify a datatype. Instead got:\n" ++ pprint info
Just x -> return x
reifyDataCon :: Name -> Q DataCon
reifyDataCon name = do
info <- reify name
case infoToDataCon info of
Nothing -> fail $ "Expected to reify a constructor. Instead got:\n" ++ pprint info
Just x -> return x
reifyDataFamily :: Name -> Q DataFamily
reifyDataFamily name = do
info <- reify name
case infoToDataFamily info of
Nothing -> fail $ "Expected to reify a data family. Instead got:\n" ++ pprint info
Just x -> return x
reifyTypeFamily :: Name -> Q TypeFamily
reifyTypeFamily name = do
info <- reify name
case infoToTypeFamily info of
Nothing -> fail $ "Expected to reify a type family. Instead got:\n" ++ pprint info
Just x -> return x
infoToDataType :: Info -> Maybe DataType
infoToDataType info = case info of
#if MIN_VERSION_template_haskell(2,11,0)
TyConI (DataD preds name tvs _kind cons _deriving) ->
#else
TyConI (DataD preds name tvs cons _deriving) ->
#endif
Just $ DataType name (map tyVarBndrName tvs) preds (concatMap conToDataCons cons)
#if MIN_VERSION_template_haskell(2,11,0)
TyConI (NewtypeD preds name tvs _kind con _deriving) ->
#else
TyConI (NewtypeD preds name tvs con _deriving) ->
#endif
Just $ DataType name (map tyVarBndrName tvs) preds (conToDataCons con)
_ -> Nothing
infoToDataFamily :: Info -> Maybe DataFamily
infoToDataFamily info = case info of
#if MIN_VERSION_template_haskell(2,11,0)
FamilyI (DataFamilyD name tvs _kind) insts ->
#else
FamilyI (FamilyD DataFam name tvs _kind) insts ->
#endif
Just $ DataFamily name (map tyVarBndrName tvs) (map go insts)
_ -> Nothing
where
#if MIN_VERSION_template_haskell(2,11,0)
go (NewtypeInstD preds name params _kind con _deriving) =
#else
go (NewtypeInstD preds name params con _deriving) =
#endif
DataInst name preds params (conToDataCons con)
#if MIN_VERSION_template_haskell(2,11,0)
go (DataInstD preds name params _kind cons _deriving) =
#else
go (DataInstD preds name params cons _deriving) =
#endif
DataInst name preds params (concatMap conToDataCons cons)
go info' = error $
"Unexpected instance in FamilyI in infoToDataInsts:\n" ++ pprint info'
infoToTypeFamily :: Info -> Maybe TypeFamily
infoToTypeFamily info = case info of
#if MIN_VERSION_template_haskell(2,11,0)
FamilyI (ClosedTypeFamilyD (TypeFamilyHead name tvs _result _injectivity) eqns) _ ->
Just $ TypeFamily name (map tyVarBndrName tvs) $ map (goEqn name) eqns
FamilyI (OpenTypeFamilyD (TypeFamilyHead name tvs _result _injectivity)) insts ->
Just $ TypeFamily name (map tyVarBndrName tvs) $ map go insts
#else
FamilyI (ClosedTypeFamilyD name tvs _kind eqns) [] ->
Just $ TypeFamily name (map tyVarBndrName tvs) $ map (goEqn name) eqns
FamilyI (FamilyD TypeFam name tvs _kind) insts ->
Just $ TypeFamily name (map tyVarBndrName tvs) $ map go insts
#endif
_ -> Nothing
where
goEqn name (TySynEqn params ty) = TypeInst name params ty
go (TySynInstD name (TySynEqn params ty)) = TypeInst name params ty
go info' = error $
"Unexpected instance in FamilyI in infoToTypeInsts:\n" ++ pprint info'
infoToDataCon :: Info -> Maybe DataCon
infoToDataCon info = case info of
#if MIN_VERSION_template_haskell(2,11,0)
DataConI name ty _parent ->
#else
DataConI name ty _parent _fixity ->
#endif
Just (typeToDataCon name ty)
_ -> Nothing
typeToDataCon :: Name -> Type -> DataCon
typeToDataCon dcName ty0 = DataCon {..}
where
(dcTvs, dcCxt, dcFields) = case ty0 of
ForallT tvs preds ty -> (map tyVarBndrName tvs, preds, typeToFields ty)
ty -> ([], [], typeToFields ty)
typeToFields = init . map (Nothing, ) . unAppsT
conToDataCons :: Con -> [DataCon]
conToDataCons = \case
NormalC name slots ->
[DataCon name [] [] (map (\(_, ty) -> (Nothing, ty)) slots)]
RecC name fields ->
[DataCon name [] [] (map (\(n, _, ty) -> (Just n, ty)) fields)]
InfixC (_, ty1) name (_, ty2) ->
[DataCon name [] [] [(Nothing, ty1), (Nothing, ty2)]]
ForallC tvs preds con ->
map (\(DataCon name tvs0 preds0 fields) ->
DataCon name (tvs0 ++ map tyVarBndrName tvs) (preds0 ++ preds) fields) (conToDataCons con)
#if MIN_VERSION_template_haskell(2,11,0)
GadtC ns slots _ ->
map (\n -> DataCon n [] [] (map (\(_, ty) -> (Nothing, ty)) slots)) ns
RecGadtC ns fields _ ->
map (\n -> DataCon n [] [] (map (\(n, _, ty) -> (Just n, ty)) fields)) ns
#endif
reifyDataTypeSubstituted :: Type -> Q DataType
reifyDataTypeSubstituted ty =
case typeToNamedCon ty of
Nothing -> fail $ "Expected a datatype, but reifyDataTypeSubstituted was applied to " ++ pprint ty
Just (n, args) -> do
dt <- reifyDataType n
let cons' = substituteTvs (M.fromList (zip (dtTvs dt) args)) (dtCons dt)
return (dt { dtCons = cons' })
substituteTvs :: Data a => M.Map Name Type -> a -> a
substituteTvs mp = transformTypes go
where
go (VarT name) | Just ty <- M.lookup name mp = ty
go ty = gmapT (substituteTvs mp) ty
transformTypes :: Data a => (Type -> Type) -> a -> a
transformTypes f = gmapT (transformTypes f) `extT` (id :: String -> String) `extT` f