singletons-th-3.2: src/Data/Singletons/TH/Util.hs
{-# LANGUAGE LambdaCase #-}
{- Data/Singletons/TH/Util.hs
(c) Richard Eisenberg 2013
rae@cs.brynmawr.edu
This file contains helper functions internal to the singletons-th package.
Users of the package should not need to consult this file.
-}
module Data.Singletons.TH.Util where
import Prelude hiding ( exp, foldl, concat, mapM, any, pred )
import Language.Haskell.TH ( pprint )
import Language.Haskell.TH.Syntax hiding ( lift )
import Language.Haskell.TH.Desugar
import Data.Char
import Control.Monad ( liftM, unless, when )
import Control.Monad.Except ( ExceptT, runExceptT, MonadError(..) )
import Control.Monad.IO.Class ( MonadIO )
import Control.Monad.Reader ( MonadReader(..), Reader, ReaderT(..) )
import Control.Monad.Trans ( MonadTrans )
import Control.Monad.Writer ( MonadWriter(..), WriterT(..), execWriterT )
import qualified Data.Map as Map
import Data.Map ( Map )
import Data.Bifunctor (second)
import Data.Foldable
import Data.Functor.Identity
import Data.Traversable
import Data.Generics
import Data.Maybe
-- like reportWarning, but generalized to any Quasi
qReportWarning :: Quasi q => String -> q ()
qReportWarning = qReport False
-- like reportError, but generalized to any Quasi
qReportError :: Quasi q => String -> q ()
qReportError = qReport True
-- | Generate a new Unique
qNewUnique :: DsMonad q => q Uniq
qNewUnique = do
Name _ flav <- qNewName "x"
case flav of
NameU n -> return n
_ -> error "Internal error: `qNewName` didn't return a NameU"
checkForRep :: Quasi q => [Name] -> q ()
checkForRep names =
when (any ((== "Rep") . nameBase) names)
(fail $ "A data type named <<Rep>> is a special case.\n" ++
"Promoting it will not work as expected.\n" ++
"Please choose another name for your data type.")
checkForRepInDecls :: Quasi q => [DDec] -> q ()
checkForRepInDecls decls =
checkForRep (allNamesIn decls)
tysOfConFields :: DConFields -> [DType]
tysOfConFields (DNormalC _ stys) = map snd stys
tysOfConFields (DRecC vstys) = map (\(_,_,ty) -> ty) vstys
recSelsOfConFields :: DConFields -> [Name]
recSelsOfConFields DNormalC{} = []
recSelsOfConFields (DRecC vstys) = map (\(n,_,_) -> n) vstys
-- Extract a data constructor's name and the number of arguments it accepts.
extractNameArgs :: DCon -> (Name, Int)
extractNameArgs (DCon _ _ n fields _) = (n, length (tysOfConFields fields))
-- Extract a data constructor's name.
extractName :: DCon -> Name
extractName (DCon _ _ n _ _) = n
-- Extract the names of a data constructor's record selectors.
extractRecSelNames :: DCon -> [Name]
extractRecSelNames (DCon _ _ _ fields _) = recSelsOfConFields fields
-- | is a valid Haskell infix data constructor (i.e., does it begin with a colon?)
isInfixDataCon :: String -> Bool
isInfixDataCon (':':_) = True
isInfixDataCon _ = False
-- | Is an identifier a legal data constructor name in Haskell? That is, is its
-- first character an uppercase letter (prefix) or a colon (infix)?
isDataConName :: Name -> Bool
isDataConName n = let first = head (nameBase n) in isUpper first || first == ':'
-- | Is an identifier uppercase?
--
-- Note that this will always return 'False' for infix names, since the concept
-- of upper- and lower-case doesn't make sense for non-alphabetic characters.
-- If you want to check if a name is legal as a data constructor, use the
-- 'isDataConName' function.
isUpcase :: Name -> Bool
isUpcase n = let first = head (nameBase n) in isUpper first
-- Make an identifier uppercase. If the identifier is infix, this acts as the
-- identity function.
upcase :: Name -> Name
upcase = mkName . toUpcaseStr noPrefix
-- make an identifier uppercase and return it as a String
toUpcaseStr :: (String, String) -- (alpha, symb) prefixes to prepend
-> Name -> String
toUpcaseStr (alpha, symb) n
| isHsLetter first
= upcase_alpha
| otherwise
= upcase_symb
where
str = nameBase n
first = head str
upcase_alpha = alpha ++ (toUpper first) : tail str
upcase_symb = symb ++ str
noPrefix :: (String, String)
noPrefix = ("", "")
-- Put an uppercase prefix on a constructor name. Takes two prefixes:
-- one for identifiers and one for symbols.
--
-- This is different from 'prefixName' in that infix constructor names always
-- start with a colon, so we must insert the prefix after the colon in order
-- for the new name to be syntactically valid.
prefixConName :: String -> String -> Name -> Name
prefixConName pre tyPre n = case (nameBase n) of
(':' : rest) -> mkName (':' : tyPre ++ rest)
alpha -> mkName (pre ++ alpha)
-- Put a prefix on a name. Takes two prefixes: one for identifiers
-- and one for symbols.
prefixName :: String -> String -> Name -> Name
prefixName pre tyPre n =
let str = nameBase n
first = head str in
if isHsLetter first
then mkName (pre ++ str)
else mkName (tyPre ++ str)
-- Put a suffix on a name. Takes two suffixes: one for identifiers
-- and one for symbols.
suffixName :: String -> String -> Name -> Name
suffixName ident symb n =
let str = nameBase n
first = head str in
if isHsLetter first
then mkName (str ++ ident)
else mkName (str ++ symb)
-- convert a number into both alphanumeric and symoblic forms
uniquePrefixes :: String -- alphanumeric prefix
-> String -- symbolic prefix
-> Uniq
-> (String, String) -- (alphanum, symbolic)
uniquePrefixes alpha symb n = (alpha ++ n_str, symb ++ convert n_str)
where
n_str = show n
convert [] = []
convert (d : ds) =
let d' = case d of
'0' -> '!'
'1' -> '#'
'2' -> '$'
'3' -> '%'
'4' -> '&'
'5' -> '*'
'6' -> '+'
'7' -> '.'
'8' -> '/'
'9' -> '>'
_ -> error "non-digit in show #"
in d' : convert ds
-- extract the kind from a TyVarBndr
extractTvbKind :: DTyVarBndr flag -> Maybe DKind
extractTvbKind (DPlainTV _ _) = Nothing
extractTvbKind (DKindedTV _ _ k) = Just k
-- extract the name from a TyVarBndr.
extractTvbName :: DTyVarBndr flag -> Name
extractTvbName (DPlainTV n _) = n
extractTvbName (DKindedTV n _ _) = n
tvbToType :: DTyVarBndr flag -> DType
tvbToType = DVarT . extractTvbName
-- If a type variable binder lacks an explicit kind, pick a default kind of
-- Type. Otherwise, leave the binder alone.
defaultTvbToTypeKind :: DTyVarBndr flag -> DTyVarBndr flag
defaultTvbToTypeKind (DPlainTV tvname f) = DKindedTV tvname f $ DConT typeKindName
defaultTvbToTypeKind tvb = tvb
-- If @Nothing@, return @Type@. If @Just k@, return @k@.
defaultMaybeToTypeKind :: Maybe DKind -> DKind
defaultMaybeToTypeKind (Just k) = k
defaultMaybeToTypeKind Nothing = DConT typeKindName
inferMaybeKindTV :: Name -> Maybe DKind -> DTyVarBndrUnit
inferMaybeKindTV n Nothing = DPlainTV n ()
inferMaybeKindTV n (Just k) = DKindedTV n () k
resultSigToMaybeKind :: DFamilyResultSig -> Maybe DKind
resultSigToMaybeKind DNoSig = Nothing
resultSigToMaybeKind (DKindSig k) = Just k
resultSigToMaybeKind (DTyVarSig DPlainTV{}) = Nothing
resultSigToMaybeKind (DTyVarSig (DKindedTV _ _ k)) = Just k
maybeKindToResultSig :: Maybe DKind -> DFamilyResultSig
maybeKindToResultSig = maybe DNoSig DKindSig
maybeSigT :: DType -> Maybe DKind -> DType
maybeSigT ty Nothing = ty
maybeSigT ty (Just ki) = ty `DSigT` ki
-- Reconstruct a vanilla function type from its individual type variable
-- binders, constraints, argument types, and result type. (See
-- Note [Vanilla-type validity checking during promotion] in
-- Data.Singletons.TH.Promote.Type for what "vanilla" means.)
ravelVanillaDType :: [DTyVarBndrSpec] -> DCxt -> [DType] -> DType -> DType
ravelVanillaDType tvbs ctxt args res =
ifNonEmpty tvbs (DForallT . DForallInvis) $
ifNonEmpty ctxt DConstrainedT $
go args
where
ifNonEmpty :: [a] -> ([a] -> b -> b) -> b -> b
ifNonEmpty [] _ z = z
ifNonEmpty l f z = f l z
go :: [DType] -> DType
go [] = res
go (h:t) = DAppT (DAppT DArrowT h) (go t)
-- Decompose a vanilla function type into its type variables, its context, its
-- argument types, and its result type. (See
-- Note [Vanilla-type validity checking during promotion] in
-- Data.Singletons.TH.Promote.Type for what "vanilla" means.)
-- If a non-vanilla construct is encountered while decomposing the function
-- type, an error is thrown monadically.
--
-- This should be contrasted with the 'unravelDType' function from
-- @th-desugar@, which supports the full gamut of function types. @singletons-th@
-- only supports a subset of these types, which is why this function is used
-- to decompose them instead.
unravelVanillaDType :: forall m. MonadFail m
=> DType -> m ([DTyVarBndrSpec], DCxt, [DType], DType)
unravelVanillaDType ty =
case unravelVanillaDType_either ty of
Left err -> fail err
Right payload -> pure payload
-- Ensures that a 'DType' is a vanilla type. (See
-- Note [Vanilla-type validity checking during promotion] in
-- Data.Singletons.TH.Promote.Type for what "vanilla" means.)
--
-- The only monadic thing that this function can do is 'fail', which it does
-- if a non-vanilla construct is encountered.
checkVanillaDType :: forall m. MonadFail m => DType -> m ()
checkVanillaDType ty =
case unravelVanillaDType_either ty of
Left err -> fail err
Right _ -> pure ()
-- The workhorse that powers unravelVanillaDType and checkVanillaDType.
-- Returns @Right payload@ upon success, and @Left error_msg@ upon failure.
unravelVanillaDType_either ::
DType -> Either String ([DTyVarBndrSpec], DCxt, [DType], DType)
unravelVanillaDType_either ty =
runIdentity $ flip runReaderT True $ runExceptT $ runUnravelM $ go_ty ty
where
go_ty :: DType -> UnravelM ([DTyVarBndrSpec], DCxt, [DType], DType)
go_ty typ = do
let (args1, res) = unravelDType typ
(args2, tvbs) <- take_tvbs args1
(args3, ctxt) <- take_ctxt args2
anons <- take_anons args3
pure (tvbs, ctxt, anons, res)
-- Process a type in a higher-order position (e.g., the @forall a. a -> a@ in
-- @(forall a. a -> a) -> b -> b@). This is only done to check for the
-- presence of higher-rank foralls or constraints, which are not permitted
-- in vanilla types.
go_higher_order_ty :: DType -> UnravelM ()
go_higher_order_ty typ = () <$ local (const False) (go_ty typ)
take_tvbs :: DFunArgs -> UnravelM (DFunArgs, [DTyVarBndrSpec])
take_tvbs (DFAForalls (DForallInvis tvbs) args) = do
rank_1 <- ask
unless rank_1 $ fail_forall "higher-rank"
_ <- traverse_ (traverse_ go_higher_order_ty . extractTvbKind) tvbs
(args', tvbs') <- take_tvbs args
pure (args', tvbs ++ tvbs')
take_tvbs (DFAForalls DForallVis{} _) = fail_vdq
take_tvbs args = pure (args, [])
take_ctxt :: DFunArgs -> UnravelM (DFunArgs, DCxt)
take_ctxt (DFACxt ctxt args) = do
rank_1 <- ask
unless rank_1 $ fail_ctxt "higher-rank"
traverse_ go_higher_order_ty ctxt
(args', ctxt') <- take_ctxt args
pure (args', ctxt ++ ctxt')
take_ctxt (DFAForalls tele _) =
case tele of
DForallInvis{} -> fail_forall "nested"
DForallVis{} -> fail_vdq
take_ctxt args = pure (args, [])
take_anons :: DFunArgs -> UnravelM [DType]
take_anons (DFAAnon anon args) = do
go_higher_order_ty anon
anons <- take_anons args
pure (anon:anons)
take_anons (DFAForalls tele _) =
case tele of
DForallInvis{} -> fail_forall "nested"
DForallVis{} -> fail_vdq
take_anons (DFACxt _ _) = fail_ctxt "nested"
take_anons DFANil = pure []
failWith :: MonadError String m => String -> m a
failWith thing = throwError $ unlines
[ "`singletons-th` does not support " ++ thing
, "In the type: " ++ pprint (sweeten ty)
]
fail_forall :: MonadError String m => String -> m a
fail_forall sort = failWith $ sort ++ " `forall`s"
fail_vdq :: MonadError String m => m a
fail_vdq = failWith "visible dependent quantification"
fail_ctxt :: MonadError String m => String -> m a
fail_ctxt sort = failWith $ sort ++ " contexts"
-- The monad that powers the internals of unravelVanillaDType_either.
--
-- * ExceptT String: records the error message upon failure.
--
-- * Reader Bool: True if we are in a rank-1 position in a type, False otherwise
newtype UnravelM a = UnravelM { runUnravelM :: ExceptT String (Reader Bool) a }
deriving (Functor, Applicative, Monad, MonadError String, MonadReader Bool)
-- count the number of arguments in a type
countArgs :: DType -> Int
countArgs ty = length $ filterDVisFunArgs args
where (args, _) = unravelDType ty
-- Collect the invisible type variable binders from a sequence of DFunArgs.
filterInvisTvbArgs :: DFunArgs -> [DTyVarBndrSpec]
filterInvisTvbArgs DFANil = []
filterInvisTvbArgs (DFACxt _ args) = filterInvisTvbArgs args
filterInvisTvbArgs (DFAAnon _ args) = filterInvisTvbArgs args
filterInvisTvbArgs (DFAForalls tele args) =
let res = filterInvisTvbArgs args in
case tele of
DForallVis _ -> res
DForallInvis tvbs' -> tvbs' ++ res
-- Infer the kind of a DTyVarBndr by using information from a DVisFunArg.
replaceTvbKind :: DVisFunArg -> DTyVarBndrUnit -> DTyVarBndrUnit
replaceTvbKind (DVisFADep tvb) _ = tvb
replaceTvbKind (DVisFAAnon k) tvb = DKindedTV (extractTvbName tvb) () k
-- changes all TyVars not to be NameU's. Workaround for GHC#11812/#17537/#19743
noExactTyVars :: Data a => a -> a
noExactTyVars = everywhere go
where
go :: Data a => a -> a
go = mkT (fix_tvb @Specificity)
`extT` fix_tvb @()
`extT` fix_ty
`extT` fix_inj_ann
fix_tvb :: Typeable flag => DTyVarBndr flag -> DTyVarBndr flag
fix_tvb (DPlainTV n f) = DPlainTV (noExactName n) f
fix_tvb (DKindedTV n f k) = DKindedTV (noExactName n) f k
fix_ty (DVarT n) = DVarT (noExactName n)
fix_ty ty = ty
fix_inj_ann (InjectivityAnn lhs rhs)
= InjectivityAnn (noExactName lhs) (map noExactName rhs)
-- changes a Name not to be a NameU. Workaround for GHC#11812/#17537/#19743
noExactName :: Name -> Name
noExactName (Name (OccName occ) (NameU unique)) = mkName (occ ++ show unique)
noExactName n = n
substKind :: Map Name DKind -> DKind -> DKind
substKind = substType
-- | Non–capture-avoiding substitution. (If you want capture-avoiding
-- substitution, use @substTy@ from "Language.Haskell.TH.Desugar.Subst".
substType :: Map Name DType -> DType -> DType
substType subst ty | Map.null subst = ty
substType subst (DForallT tele inner_ty)
= DForallT tele' inner_ty'
where
(subst', tele') = subst_tele subst tele
inner_ty' = substType subst' inner_ty
substType subst (DConstrainedT cxt inner_ty) =
DConstrainedT (map (substType subst) cxt) (substType subst inner_ty)
substType subst (DAppT ty1 ty2) = substType subst ty1 `DAppT` substType subst ty2
substType subst (DAppKindT ty ki) = substType subst ty `DAppKindT` substType subst ki
substType subst (DSigT ty ki) = substType subst ty `DSigT` substType subst ki
substType subst (DVarT n) =
case Map.lookup n subst of
Just ki -> ki
Nothing -> DVarT n
substType _ ty@(DConT {}) = ty
substType _ ty@(DArrowT) = ty
substType _ ty@(DLitT {}) = ty
substType _ ty@DWildCardT = ty
subst_tele :: Map Name DKind -> DForallTelescope -> (Map Name DKind, DForallTelescope)
subst_tele s (DForallInvis tvbs) = second DForallInvis $ subst_tvbs s tvbs
subst_tele s (DForallVis tvbs) = second DForallVis $ subst_tvbs s tvbs
subst_tvbs :: Map Name DKind -> [DTyVarBndr flag] -> (Map Name DKind, [DTyVarBndr flag])
subst_tvbs = mapAccumL subst_tvb
subst_tvb :: Map Name DKind -> DTyVarBndr flag -> (Map Name DKind, DTyVarBndr flag)
subst_tvb s tvb@(DPlainTV n _) = (Map.delete n s, tvb)
subst_tvb s (DKindedTV n f k) = (Map.delete n s, DKindedTV n f (substKind s k))
dropTvbKind :: DTyVarBndr flag -> DTyVarBndr flag
dropTvbKind tvb@(DPlainTV {}) = tvb
dropTvbKind (DKindedTV n f _) = DPlainTV n f
-- apply a type to a list of types
foldType :: DType -> [DType] -> DType
foldType = foldl DAppT
-- apply a type to a list of type variable binders
foldTypeTvbs :: DType -> [DTyVarBndr flag] -> DType
foldTypeTvbs ty = foldType ty . map tvbToType
-- Construct a data type's variable binders, possibly using fresh variables
-- from the data type's kind signature.
buildDataDTvbs :: DsMonad q => [DTyVarBndrUnit] -> Maybe DKind -> q [DTyVarBndrUnit]
buildDataDTvbs tvbs mk = do
extra_tvbs <- mkExtraDKindBinders $ fromMaybe (DConT typeKindName) mk
pure $ tvbs ++ extra_tvbs
-- apply an expression to a list of expressions
foldExp :: DExp -> [DExp] -> DExp
foldExp = foldl DAppE
-- choose the first non-empty list
orIfEmpty :: [a] -> [a] -> [a]
orIfEmpty [] x = x
orIfEmpty x _ = x
-- build a pattern match over several expressions, each with only one pattern
multiCase :: [DExp] -> [DPat] -> DExp -> DExp
multiCase [] [] body = body
multiCase scruts pats body =
DCaseE (mkTupleDExp scruts) [DMatch (mkTupleDPat pats) body]
-- a monad transformer for writing a monoid alongside returning a Q
newtype QWithAux m q a = QWA { runQWA :: WriterT m q a }
deriving ( Functor, Applicative, Monad, MonadTrans
, MonadWriter m, MonadReader r
, MonadFail, MonadIO, Quasi, DsMonad )
-- run a computation with an auxiliary monoid, discarding the monoid result
evalWithoutAux :: Quasi q => QWithAux m q a -> q a
evalWithoutAux = liftM fst . runWriterT . runQWA
-- run a computation with an auxiliary monoid, returning only the monoid result
evalForAux :: Quasi q => QWithAux m q a -> q m
evalForAux = execWriterT . runQWA
-- run a computation with an auxiliary monoid, return both the result
-- of the computation and the monoid result
evalForPair :: QWithAux m q a -> q (a, m)
evalForPair = runWriterT . runQWA
-- in a computation with an auxiliary map, add a binding to the map
addBinding :: (Quasi q, Ord k) => k -> v -> QWithAux (Map.Map k v) q ()
addBinding k v = tell (Map.singleton k v)
-- in a computation with an auxiliar list, add an element to the list
addElement :: Quasi q => elt -> QWithAux [elt] q ()
addElement elt = tell [elt]
-- | Call 'lookupTypeNameWithLocals' first to ensure we have a 'Name' in the
-- type namespace, then call 'dsReify'.
-- See also Note [Using dsReifyTypeNameInfo when promoting instances]
-- in Data.Singletons.TH.Promote.
dsReifyTypeNameInfo :: DsMonad q => Name -> q (Maybe DInfo)
dsReifyTypeNameInfo ty_name = do
mb_name <- lookupTypeNameWithLocals (nameBase ty_name)
case mb_name of
Just n -> dsReify n
Nothing -> pure Nothing
-- lift concatMap into a monad
-- could this be more efficient?
concatMapM :: (Monad monad, Monoid monoid, Traversable t)
=> (a -> monad monoid) -> t a -> monad monoid
concatMapM fn list = do
bss <- mapM fn list
return $ fold bss
-- like GHC's
mapMaybeM :: Monad m => (a -> m (Maybe b)) -> [a] -> m [b]
mapMaybeM _ [] = return []
mapMaybeM f (x:xs) = do
y <- f x
ys <- mapMaybeM f xs
return $ case y of
Nothing -> ys
Just z -> z : ys
-- make a one-element list
listify :: a -> [a]
listify = (:[])
fstOf3 :: (a,b,c) -> a
fstOf3 (a,_,_) = a
liftFst :: (a -> b) -> (a, c) -> (b, c)
liftFst f (a, c) = (f a, c)
liftSnd :: (a -> b) -> (c, a) -> (c, b)
liftSnd f (c, a) = (c, f a)
snocView :: [a] -> ([a], a)
snocView [] = error "snocView nil"
snocView [x] = ([], x)
snocView (x : xs) = liftFst (x:) (snocView xs)
partitionWith :: (a -> Either b c) -> [a] -> ([b], [c])
partitionWith f = go [] []
where go bs cs [] = (reverse bs, reverse cs)
go bs cs (a:as) =
case f a of
Left b -> go (b:bs) cs as
Right c -> go bs (c:cs) as
partitionWithM :: Monad m => (a -> m (Either b c)) -> [a] -> m ([b], [c])
partitionWithM f = go [] []
where go bs cs [] = return (reverse bs, reverse cs)
go bs cs (a:as) = do
fa <- f a
case fa of
Left b -> go (b:bs) cs as
Right c -> go bs (c:cs) as
partitionLetDecs :: [DDec] -> ([DLetDec], [DDec])
partitionLetDecs = partitionWith (\case DLetDec ld -> Left ld
dec -> Right dec)
{-# INLINEABLE zipWith3M #-}
zipWith3M :: Monad m => (a -> b -> m c) -> [a] -> [b] -> m [c]
zipWith3M f (a:as) (b:bs) = (:) <$> f a b <*> zipWith3M f as bs
zipWith3M _ _ _ = return []
mapAndUnzip3M :: Monad m => (a -> m (b,c,d)) -> [a] -> m ([b],[c],[d])
mapAndUnzip3M _ [] = return ([],[],[])
mapAndUnzip3M f (x:xs) = do
(r1, r2, r3) <- f x
(rs1, rs2, rs3) <- mapAndUnzip3M f xs
return (r1:rs1, r2:rs2, r3:rs3)
-- is it a letter or underscore?
isHsLetter :: Char -> Bool
isHsLetter c = isLetter c || c == '_'