{-# LANGUAGE CPP #-}
{-# LANGUAGE OverloadedStrings #-}
module Wingman.GHC where
import Control.Monad.State
import Control.Monad.Trans.Maybe (MaybeT(..))
import Data.Bool (bool)
import Data.Coerce (coerce)
import Data.Function (on)
import Data.Functor ((<&>))
import Data.List (isPrefixOf)
import qualified Data.Map as M
import Data.Maybe (isJust)
import Data.Set (Set)
import qualified Data.Set as S
import Data.Traversable
import Development.IDE.GHC.Compat
import Development.IDE.GHC.Compat.Util
import GHC.SourceGen (lambda)
import Generics.SYB (Data, everything, everywhere, listify, mkQ, mkT)
import Wingman.StaticPlugin (pattern MetaprogramSyntax)
import Wingman.Types
#if __GLASGOW_HASKELL__ >= 900
import GHC.Tc.Utils.TcType
#endif
tcTyVar_maybe :: Type -> Maybe Var
tcTyVar_maybe ty | Just ty' <- tcView ty = tcTyVar_maybe ty'
tcTyVar_maybe (CastTy ty _) = tcTyVar_maybe ty -- look through casts, as
-- this is only used for
-- e.g., FlexibleContexts
tcTyVar_maybe (TyVarTy v) = Just v
tcTyVar_maybe _ = Nothing
instantiateType :: Type -> ([TyVar], Type)
instantiateType t = do
let vs = tyCoVarsOfTypeList t
vs' = fmap cloneTyVar vs
subst = foldr (\(v,t) a -> extendTCvSubst a v $ TyVarTy t) emptyTCvSubst
$ zip vs vs'
in (vs', substTy subst t)
cloneTyVar :: TyVar -> TyVar
cloneTyVar t =
let uniq = getUnique t
some_magic_char = 'w' -- 'w' for wingman ;D
in setVarUnique t $ newTagUnique uniq some_magic_char
------------------------------------------------------------------------------
-- | Is this a function type?
isFunction :: Type -> Bool
isFunction (tacticsSplitFunTy -> (_, _, [], _)) = False
isFunction _ = True
------------------------------------------------------------------------------
-- | Split a function, also splitting out its quantified variables and theta
-- context.
tacticsSplitFunTy :: Type -> ([TyVar], ThetaType, [Scaled Type], Type)
tacticsSplitFunTy t
= let (vars, theta, t') = tcSplitNestedSigmaTys t
(args, res) = tcSplitFunTys t'
in (vars, theta, args, res)
------------------------------------------------------------------------------
-- | Rip the theta context out of a regular type.
tacticsThetaTy :: Type -> ThetaType
tacticsThetaTy (tcSplitSigmaTy -> (_, theta, _)) = theta
------------------------------------------------------------------------------
-- | Get the data cons of a type, if it has any.
tacticsGetDataCons :: Type -> Maybe ([DataCon], [Type])
tacticsGetDataCons ty
| Just (_, ty') <- tcSplitForAllTyVarBinder_maybe ty
= tacticsGetDataCons ty'
tacticsGetDataCons ty
| Just _ <- algebraicTyCon ty
= splitTyConApp_maybe ty <&> \(tc, apps) ->
( filter (not . dataConCannotMatch apps) $ tyConDataCons tc
, apps
)
tacticsGetDataCons _ = Nothing
------------------------------------------------------------------------------
-- | Instantiate all of the quantified type variables in a type with fresh
-- skolems.
freshTyvars :: MonadState TacticState m => Type -> m Type
freshTyvars t = do
let (tvs, _, _, _) = tacticsSplitFunTy t
reps <- fmap M.fromList
$ for tvs $ \tv -> do
uniq <- freshUnique
pure (tv, setTyVarUnique tv uniq)
pure $
everywhere
(mkT $ \tv -> M.findWithDefault tv tv reps
) $ snd $ tcSplitForAllTyVars t
------------------------------------------------------------------------------
-- | Given a datacon, extract its record fields' names and types. Returns
-- nothing if the datacon is not a record.
getRecordFields :: ConLike -> Maybe [(OccName, CType)]
getRecordFields dc =
case conLikeFieldLabels dc of
[] -> Nothing
lbls -> for lbls $ \lbl -> do
let ty = conLikeFieldType dc $ flLabel lbl
pure (mkVarOccFS $ flLabel lbl, CType ty)
------------------------------------------------------------------------------
-- | Is this an algebraic type?
algebraicTyCon :: Type -> Maybe TyCon
algebraicTyCon ty
| Just (_, ty') <- tcSplitForAllTyVarBinder_maybe ty
= algebraicTyCon ty'
algebraicTyCon (splitTyConApp_maybe -> Just (tycon, _))
| tycon == intTyCon = Nothing
| tycon == floatTyCon = Nothing
| tycon == doubleTyCon = Nothing
| tycon == charTyCon = Nothing
| tycon == funTyCon = Nothing
| otherwise = Just tycon
algebraicTyCon _ = Nothing
------------------------------------------------------------------------------
-- | We can't compare 'RdrName' for equality directly. Instead, sloppily
-- compare them by their 'OccName's.
eqRdrName :: RdrName -> RdrName -> Bool
eqRdrName = (==) `on` occNameString . occName
------------------------------------------------------------------------------
-- | Compare two 'OccName's for unqualified equality.
sloppyEqOccName :: OccName -> OccName -> Bool
sloppyEqOccName = (==) `on` occNameString
------------------------------------------------------------------------------
-- | Does this thing contain any references to 'HsVar's with the given
-- 'RdrName'?
containsHsVar :: Data a => RdrName -> a -> Bool
containsHsVar name x = not $ null $ listify (
\case
((HsVar _ (L _ a)) :: HsExpr GhcPs) | eqRdrName a name -> True
_ -> False
) x
------------------------------------------------------------------------------
-- | Does this thing contain any holes?
containsHole :: Data a => a -> Bool
containsHole x = not $ null $ listify (
\case
((HsVar _ (L _ name)) :: HsExpr GhcPs) -> isHole $ occName name
MetaprogramSyntax _ -> True
_ -> False
) x
------------------------------------------------------------------------------
-- | Check if an 'OccName' is a hole
isHole :: OccName -> Bool
-- TODO(sandy): Make this more robust
isHole = isPrefixOf "_" . occNameString
------------------------------------------------------------------------------
-- | Get all of the referenced occnames.
allOccNames :: Data a => a -> Set OccName
allOccNames = everything (<>) $ mkQ mempty $ \case
a -> S.singleton a
------------------------------------------------------------------------------
-- | Unpack the relevant parts of a 'Match'
#if __GLASGOW_HASKELL__ >= 900
pattern AMatch :: HsMatchContext (NoGhcTc GhcPs) -> [Pat GhcPs] -> HsExpr GhcPs -> Match GhcPs (LHsExpr GhcPs)
#else
pattern AMatch :: HsMatchContext (NameOrRdrName (IdP GhcPs)) -> [Pat GhcPs] -> HsExpr GhcPs -> Match GhcPs (LHsExpr GhcPs)
#endif
pattern AMatch ctx pats body <-
Match { m_ctxt = ctx
, m_pats = fmap fromPatCompat -> pats
, m_grhss = UnguardedRHSs (unLoc -> body)
}
pattern SingleLet :: IdP GhcPs -> [Pat GhcPs] -> HsExpr GhcPs -> HsExpr GhcPs -> HsExpr GhcPs
pattern SingleLet bind pats val expr <-
HsLet _
(HsValBinds _
(ValBinds _ (bagToList ->
[L _ (FunBind {fun_id = (L _ bind), fun_matches = (MG _ (L _ [L _ (AMatch _ pats val)]) _)})]) _))
(L _ expr)
------------------------------------------------------------------------------
-- | A pattern over the otherwise (extremely) messy AST for lambdas.
pattern Lambda :: [Pat GhcPs] -> HsExpr GhcPs -> HsExpr GhcPs
pattern Lambda pats body <-
HsLam _
MG {mg_alts = L _ [L _ (AMatch _ pats body) ]}
where
-- If there are no patterns to bind, just stick in the body
Lambda [] body = body
Lambda pats body = lambda pats body
------------------------------------------------------------------------------
-- | A GRHS that caontains no guards.
pattern UnguardedRHSs :: LHsExpr p -> GRHSs p (LHsExpr p)
pattern UnguardedRHSs body <-
GRHSs {grhssGRHSs = [L _ (GRHS _ [] body)]}
------------------------------------------------------------------------------
-- | A match with a single pattern. Case matches are always 'SinglePatMatch'es.
pattern SinglePatMatch :: PatCompattable p => Pat p -> LHsExpr p -> Match p (LHsExpr p)
pattern SinglePatMatch pat body <-
Match { m_pats = [fromPatCompat -> pat]
, m_grhss = UnguardedRHSs body
}
------------------------------------------------------------------------------
-- | Helper function for defining the 'Case' pattern.
unpackMatches :: PatCompattable p => [Match p (LHsExpr p)] -> Maybe [(Pat p, LHsExpr p)]
unpackMatches [] = Just []
unpackMatches (SinglePatMatch pat body : matches) =
((pat, body):) <$> unpackMatches matches
unpackMatches _ = Nothing
------------------------------------------------------------------------------
-- | A pattern over the otherwise (extremely) messy AST for lambdas.
pattern Case :: PatCompattable p => HsExpr p -> [(Pat p, LHsExpr p)] -> HsExpr p
pattern Case scrutinee matches <-
HsCase _ (L _ scrutinee)
MG {mg_alts = L _ (fmap unLoc -> unpackMatches -> Just matches)}
------------------------------------------------------------------------------
-- | Like 'Case', but for lambda cases.
pattern LamCase :: PatCompattable p => [(Pat p, LHsExpr p)] -> HsExpr p
pattern LamCase matches <-
HsLamCase _
MG {mg_alts = L _ (fmap unLoc -> unpackMatches -> Just matches)}
------------------------------------------------------------------------------
-- | Can ths type be lambda-cased?
--
-- Return: 'Nothing' if no
-- @Just False@ if it can't be homomorphic
-- @Just True@ if it can
lambdaCaseable :: Type -> Maybe Bool
#if __GLASGOW_HASKELL__ >= 900
lambdaCaseable (splitFunTy_maybe -> Just (_multiplicity, arg, res))
#else
lambdaCaseable (splitFunTy_maybe -> Just (arg, res))
#endif
| isJust (algebraicTyCon arg)
= Just $ isJust $ algebraicTyCon res
lambdaCaseable _ = Nothing
class PatCompattable p where
fromPatCompat :: PatCompat p -> Pat p
toPatCompat :: Pat p -> PatCompat p
#if __GLASGOW_HASKELL__ == 808
instance PatCompattable GhcTc where
fromPatCompat = id
toPatCompat = id
instance PatCompattable GhcPs where
fromPatCompat = id
toPatCompat = id
type PatCompat pass = Pat pass
#else
instance PatCompattable GhcTc where
fromPatCompat = unLoc
toPatCompat = noLoc
instance PatCompattable GhcPs where
fromPatCompat = unLoc
toPatCompat = noLoc
type PatCompat pass = LPat pass
#endif
------------------------------------------------------------------------------
-- | Should make sure it's a fun bind
pattern TopLevelRHS
:: OccName
-> [PatCompat GhcTc]
-> LHsExpr GhcTc
-> HsLocalBindsLR GhcTc GhcTc
-> Match GhcTc (LHsExpr GhcTc)
pattern TopLevelRHS name ps body where_binds <-
Match _
(FunRhs (L _ (occName -> name)) _ _)
ps
(GRHSs _
[L _ (GRHS _ [] body)] (L _ where_binds))
------------------------------------------------------------------------------
-- | In GHC 8.8, sometimes patterns are wrapped in 'XPat'.
-- The nitty gritty details are explained at
-- https://blog.shaynefletcher.org/2020/03/ghc-haskell-pats-and-lpats.html
--
-- We need to remove these in order to succesfull find patterns.
unXPat :: Pat GhcPs -> Pat GhcPs
#if __GLASGOW_HASKELL__ == 808
unXPat (XPat (L _ pat)) = unXPat pat
#endif
unXPat pat = pat
liftMaybe :: Monad m => Maybe a -> MaybeT m a
liftMaybe a = MaybeT $ pure a
------------------------------------------------------------------------------
-- | Get the type of an @HsExpr GhcTc@. This is slow and you should prefer to
-- not use it, but sometimes it can't be helped.
typeCheck :: HscEnv -> TcGblEnv -> HsExpr GhcTc -> IO (Maybe Type)
typeCheck hscenv tcg = fmap snd . initDs hscenv tcg . fmap exprType . dsExpr
------------------------------------------------------------------------------
-- | Expand type and data families
normalizeType :: Context -> Type -> Type
normalizeType ctx ty =
let ty' = expandTyFam ctx ty
in case tcSplitTyConApp_maybe ty' of
Just (tc, tys) ->
-- try to expand any data families
case tcLookupDataFamInst_maybe (ctxFamInstEnvs ctx) tc tys of
Just (dtc, dtys, _) -> mkAppTys (mkTyConTy dtc) dtys
Nothing -> ty'
Nothing -> ty'
------------------------------------------------------------------------------
-- | Expand type families
expandTyFam :: Context -> Type -> Type
expandTyFam ctx = snd . normaliseType (ctxFamInstEnvs ctx) Nominal
------------------------------------------------------------------------------
-- | Like 'tcUnifyTy', but takes a list of skolems to prevent unification of.
tryUnifyUnivarsButNotSkolems :: Set TyVar -> CType -> CType -> Maybe TCvSubst
tryUnifyUnivarsButNotSkolems skolems goal inst =
tryUnifyUnivarsButNotSkolemsMany skolems $ coerce [(goal, inst)]
------------------------------------------------------------------------------
-- | Like 'tryUnifyUnivarsButNotSkolems', but takes a list
-- of pairs of types to unify.
tryUnifyUnivarsButNotSkolemsMany :: Set TyVar -> [(Type, Type)] -> Maybe TCvSubst
tryUnifyUnivarsButNotSkolemsMany skolems (unzip -> (goal, inst)) =
tcUnifyTys
(bool BindMe Skolem . flip S.member skolems)
inst
goal
updateSubst :: TCvSubst -> TacticState -> TacticState
updateSubst subst s = s { ts_unifier = unionTCvSubst subst (ts_unifier s) }
------------------------------------------------------------------------------
-- | Get the class methods of a 'PredType', correctly dealing with
-- instantiation of quantified class types.
methodHypothesis :: PredType -> Maybe [HyInfo CType]
methodHypothesis ty = do
(tc, apps) <- splitTyConApp_maybe ty
cls <- tyConClass_maybe tc
let methods = classMethods cls
tvs = classTyVars cls
subst = zipTvSubst tvs apps
pure $ methods <&> \method ->
let (_, _, ty) = tcSplitSigmaTy $ idType method
in ( HyInfo (occName method) (ClassMethodPrv $ Uniquely cls) $ CType $ substTy subst ty
)