{-# LANGUAGE CPP, ConstraintKinds, ImplicitParams, TemplateHaskell #-} -- | Template Haskell script to promote a type family to first class. module Fcf.Family.TH ( -- * Generate boilerplate fcfify , fcfifySkip , fcfify' -- * Using promoted families , promoteFamily , promoteNDFamily , familyName , applyFamily , consTuple , paramsProxy -- * Predicates , isTypeFamily , isTypeSynonym , isTypeFamilyOrSynonym ) where import Control.Applicative (liftA2) import Control.Monad (when) import Data.Function (on) import Data.Functor (($>)) import Data.List (sort) import Data.Maybe (fromMaybe) import Data.Foldable (foldl') import Data.Traversable (for) import Data.Set (Set) import qualified Data.Set as Set import Language.Haskell.TH import Language.Haskell.TH.Syntax (getQ, putQ) import Fcf.Core import Fcf.Family hiding (Name) -- | Generate the boilerplate needed to promote a type family to first class. -- -- Required extensions: -- -- - @DataKinds@ -- - @PolyKinds@ -- - @TypeFamilies@ -- -- If 'fcfify' is called more than once with the same 'Name' in the same module, -- only the first invocation generates declarations; subsequent declarations -- return the empty list, avoiding duplicate declarations in the current module. -- -- For a stateless variant, use 'fcfify''. -- -- See "Fcf.Family" for details on the encoding. fcfify :: Name -> Q [Dec] fcfify name = do check <- checkFcfified name if check then pure [] else fcfify' name -- | Mark a type family as already fcifified. fcfifySkip :: Name -> Q [Dec] fcfifySkip name = checkFcfified name $> [] -- | Store invocations of 'fcfify' to avoid generating duplicate instances -- in the current module (a minor performance optimization). newtype Fcfified = Fcfified (Set Name) -- | Check whether we've already seen this name. -- Add the name to the registered set. checkFcfified :: Name -> Q Bool checkFcfified name = do Fcfified seen <- fromMaybe (Fcfified Set.empty) <$> getQ let check = name `Set.member` seen when (not check) (putQ (Fcfified (Set.insert name seen))) pure check -- | Generate the boilerplate needed to promote a type family to first class. -- -- Unlike 'fcfify', this always returns the same declarations for the same -- named type. fcfify' :: Name -> Q [Dec] fcfify' name = reifyTyInfo name >>= fcfifyInfo where ?funName = "fcfify" ?name = name -- | Get the quoted fcf 'Fcf.Core.Family.Name' of an existing type family. familyName :: Name -> Type familyName name = PromotedT 'MkName `AppT` lit (fromMaybe "" (namePackage name)) `AppT` lit (fromMaybe "" (nameModule name)) `AppT` lit (nameBase name) where lit = LitT . StrTyLit -- | Promote a fcfified family, returning its partially applied 'Family' and -- its arity. The result can be applied to a 'consTuple' of the appropriate size, promoteFamily :: Name -> Q (Type, Int) promoteFamily = promoteFamily_ ''Family -- | Promote a fcfified family, returning its partially applied 'Family' and -- its arity. The result can be applied to a 'consTuple' of the appropriate size, promoteNDFamily :: Name -> Q (Type, Int) promoteNDFamily = promoteFamily_ ''NDFamily promoteFamily_ :: Name -> Name -> Q (Type, Int) promoteFamily_ _Family name = do info <- reifyTyInfo name let arity = length (tiArgs info) pure (ConT _Family `AppT` tiNameT info `AppT` paramsProxy' info, arity) where ?funName = "promoteFamily_" -- | Apply a promoted family. -- -- If there are more arguments than the arity of the family (as returned by 'promoteFamily'), -- they are split and applied properly: -- the family's main arguments are collected in a 'consTuple' and -- the rest are applied with 'AppT'. -- -- If there are fewer arguments than the arity, the result is nonsense. applyFamily :: Name -> [Q Type] -> Q Type applyFamily name argsQ = do (fam, arity) <- promoteFamily name (args1, args2) <- splitAt arity <$> sequenceA argsQ pure (fam `AppT` consTuple args1 `appsT` args2) paramsProxy :: Name -> Q Type paramsProxy name = paramsProxy' <$> reifyTyInfo name where ?funName = "paramsProxy" paramsProxy' :: TyInfo -> Type paramsProxy' info = go (length (tiParams info)) where go 0 = ConT ''P0 go n = ConT ''PS `AppT` go (n-1) reifyTyInfo :: (?funName :: String) => Name -> Q TyInfo reifyTyInfo name = do let ?name = name info <- reify name case info of FamilyI dec _ -> reifyTyInfoDec dec TyConI dec -> reifyTyInfoDec dec _ -> errorNotType -- | 'True' if it is a type family (open or closed). isTypeFamily :: Name -> Q Bool isTypeFamily name = isTypeFamilyInfo <$> reify name -- | 'True' if it is a type synonym. isTypeSynonym :: Name -> Q Bool isTypeSynonym name = isTypeSynonymInfo <$> reify name -- | 'True' if it is a type family or synonym. isTypeFamilyOrSynonym :: Name -> Q Bool isTypeFamilyOrSynonym name = liftA2 (||) isTypeFamilyInfo isTypeSynonymInfo <$> reify name isTypeFamilyInfo :: Info -> Bool isTypeFamilyInfo (FamilyI (OpenTypeFamilyD _) _) = True isTypeFamilyInfo (FamilyI (ClosedTypeFamilyD _ _) _) = True isTypeFamilyInfo _ = False isTypeSynonymInfo :: Info -> Bool isTypeSynonymInfo (TyConI (TySynD _ _ _)) = True isTypeSynonymInfo _ = False -- type ErrCtxt = (?funName :: String, ?name :: Name) errorNotType :: ErrCtxt => Q a errorNotType = fail (?funName ++ ": unexpected name, " ++ show ?name ++ " is not a type family or type synonym.") -- Example: -- -- @ -- -- Input -- type F a b c = (...) -- -- -- Output -- type instance Params F -- @ reifyTyInfoDec :: ErrCtxt => Dec -> Q TyInfo reifyTyInfoDec (TySynD name args _) = mkInfoHead name args StarT -- TODO: don't assume result kind is Type reifyTyInfoDec (OpenTypeFamilyD t) = reifyTyInfoTFH t reifyTyInfoDec (ClosedTypeFamilyD t _) = reifyTyInfoTFH t reifyTyInfoDec _ = errorNotType reifyTyInfoTFH :: ErrCtxt => TypeFamilyHead -> Q TyInfo reifyTyInfoTFH (TypeFamilyHead name args resSig _) = do res <- getRes resSig mkInfoHead name args res getRes :: ErrCtxt => FamilyResultSig -> Q Type getRes NoSig = fail (?funName ++ ": implicit result type not supported") getRes (KindSig k) = pure k getRes (TyVarSig (KindedTV _ _ k)) = pure k getRes (TyVarSig PlainTV{}) = fail (?funName ++ ": implicit result type not supported") -- mkInfoHead :: Name -> [TyVarBndr a] -> Type -> Q TyInfo mkInfoHead name args res = do args' <- for args (\arg -> case arg of PlainTV _ _ -> fail "unexpected unnanotated arguments" -- as far as I understand, the binders given by reify are always annotated so this shouldn't happen KindedTV v _ k -> pure (v, k)) let params = collectParams args' res pure (mkTyInfo name params args' res) collectParams :: [(Name, Type)] -> Type -> [Name] collectParams args res = collect Set.empty args where collect bound [] = snd (addVars bound [] (getVars res)) -- collect parameters from the result type collect bound ((v, k) : vs) = let (bound', ws) = addVars bound [] (getVars k) in ws ++ collect (Set.insert v bound') vs addVars bound ws [] = (bound, reverse ws) addVars bound ws (x : xs) | Set.member x bound = addVars bound ws xs | otherwise = addVars (Set.insert x bound) (x : ws) xs data TyInfo = TyInfo { tiName :: Name , tiNameT :: Type -- ^ Encoding of name as a 'Name' , tiParams :: [Name] , tiParamsT :: Type -- ^ Params as a tuple , tiArgs :: [(Name, Type)] , tiArgsT :: Type , tiRes :: Type } appsT :: Type -> [Type] -> Type appsT = foldl' AppT mkTyInfo :: Name -> [Name] -> [(Name, Type)] -> Type -> TyInfo mkTyInfo name params args res = TyInfo { tiName = name , tiNameT = familyName name , tiParams = params , tiParamsT = consTuple (VarT <$> params) , tiArgs = args , tiArgsT = consTuple (uncurry (SigT . VarT) <$> args) , tiRes = res } -- | Construct a tuple suitable for a 'Family' argument. consTuple :: [Type] -> Type consTuple = consTuple_ (PromotedTupleT 2) (PromotedTupleT 0) consTupleT :: [Type] -> Type consTupleT = consTuple_ (TupleT 2) (TupleT 0) consTuple_ :: Type -> Type -> [Type] -> Type consTuple_ _ unit [] = unit consTuple_ tup unit (t : ts) = tup `AppT` t `AppT` consTuple_ tup unit ts -- fcfifyInfo :: ErrCtxt => TyInfo -> Q [Dec] fcfifyInfo info = do paramsD <- declareParams info argsD <- declareArgs info resD <- declareRes info familyD <- declareFamily info pure [paramsD, argsD, resD, familyD] getVars :: Type -> [Name] getVars (VarT v) = [v] getVars (AppT t t') = getVars t ++ getVars t' getVars (AppKindT t t') = getVars t ++ getVars t' getVars (SigT t k) = getVars t ++ getVars k getVars (InfixT t _ t') = getVars t ++ getVars t' getVars (UInfixT t _ t') = getVars t ++ getVars t' getVars (ParensT t) = getVars t #if MIN_VERSION_template_haskell(2,19,0) getVars (PromotedInfixT t _ t') = getVars t ++ getVars t' getVars (PromotedUInfixT t _ t') = getVars t ++ getVars t' #endif getVars _ = [] declareParams :: TyInfo -> Q Dec declareParams info = do let nParams = length (tiParams info) pure (TySynInstD (TySynEqn Nothing (ConT ''Params `AppT` tiNameT info) (consTupleT (replicate nParams StarT)))) -- TODO: don't guess Type for all params declareArgs :: TyInfo -> Q Dec declareArgs info = do pure (TySynInstD (TySynEqn Nothing (ConT ''Args_ `AppT` tiNameT info `AppT` tiParamsT info) (consTupleT (snd <$> tiArgs info)))) declareRes :: TyInfo -> Q Dec declareRes info = do pure (TySynInstD (TySynEqn Nothing (ConT ''Res_ `AppT` tiNameT info `AppT` tiParamsT info `AppT` if isDT info then tiArgsT info else WildCardT) (tiRes info))) isDT :: TyInfo -> Bool isDT info = not (null (intersection (fst <$> tiArgs info) (getVars (tiRes info)))) intersection :: Ord a => [a] -> [a] -> [a] intersection = intersectionSorted `on` sort intersectionSorted :: Ord a => [a] -> [a] -> [a] intersectionSorted [] _ = [] intersectionSorted _ [] = [] intersectionSorted xxs@(x : xs) yys@(y : ys) = case compare x y of EQ -> x : intersectionSorted xs ys LT -> intersectionSorted xs yys GT -> intersectionSorted xxs ys declareFamily :: TyInfo -> Q Dec declareFamily info = do pure (TySynInstD (TySynEqn Nothing (ConT ''Eval `AppT` (ConT ''Family_ `AppT` tiNameT info `AppT` SigT WildCardT (WildCardT `AppT` tiParamsT info) `AppT` tiArgsT info)) (foldl' AppT (ConT (tiName info)) (VarT . fst <$> tiArgs info))))