module GHC.Iface.Ext.Utils where
import GHC.Prelude
import GHC.Core.Map
import GHC.Driver.Session ( DynFlags )
import GHC.Data.FastString ( FastString, mkFastString )
import GHC.Iface.Type
import GHC.Core.Multiplicity
import GHC.Types.Name hiding (varName)
import GHC.Types.Name.Set
import GHC.Utils.Outputable hiding ( (<>) )
import qualified GHC.Utils.Outputable as O
import GHC.Types.SrcLoc
import GHC.CoreToIface
import GHC.Core.TyCon
import GHC.Core.TyCo.Rep
import GHC.Core.Type
import GHC.Types.Var
import GHC.Types.Var.Env
import GHC.Iface.Ext.Types
import qualified Data.Map as M
import qualified Data.Set as S
import qualified Data.IntMap.Strict as IM
import qualified Data.Array as A
import Data.Data ( typeOf, typeRepTyCon, Data(toConstr) )
import Data.Maybe ( maybeToList, mapMaybe)
import Data.Monoid
import Data.List (find)
import Data.Traversable ( for )
import Data.Coerce
import Control.Monad.Trans.State.Strict hiding (get)
import Control.Monad.Trans.Reader
import qualified Data.Tree as Tree
type RefMap a = M.Map Identifier [(Span, IdentifierDetails a)]
generateReferencesMap
:: Foldable f
=> f (HieAST a)
-> RefMap a
generateReferencesMap = foldr (\ast m -> M.unionWith (++) (go ast) m) M.empty
where
go ast = M.unionsWith (++) (this : map go (nodeChildren ast))
where
this = fmap (pure . (nodeSpan ast,)) $ sourcedNodeIdents $ sourcedNodeInfo ast
renderHieType :: DynFlags -> HieTypeFix -> String
renderHieType dflags ht = renderWithStyle (initSDocContext dflags defaultUserStyle) (ppr $ hieTypeToIface ht)
resolveVisibility :: Type -> [Type] -> [(Bool,Type)]
resolveVisibility kind ty_args
= go (mkEmptyTCvSubst in_scope) kind ty_args
where
in_scope = mkInScopeSet (tyCoVarsOfTypes ty_args)
go _ _ [] = []
go env ty ts
| Just ty' <- coreView ty
= go env ty' ts
go env (ForAllTy (Bndr tv vis) res) (t:ts)
| isVisibleArgFlag vis = (True , t) : ts'
| otherwise = (False, t) : ts'
where
ts' = go (extendTvSubst env tv t) res ts
go env (FunTy { ft_res = res }) (t:ts)
= (True,t) : (go env res ts)
go env (TyVarTy tv) ts
| Just ki <- lookupTyVar env tv = go env ki ts
go env kind (t:ts) = (True, t) : (go env kind ts)
foldType :: (HieType a -> a) -> HieTypeFix -> a
foldType f (Roll t) = f $ fmap (foldType f) t
selectPoint :: HieFile -> (Int,Int) -> Maybe (HieAST Int)
selectPoint hf (sl,sc) = getFirst $
flip foldMap (M.toList (getAsts $ hie_asts hf)) $ \(fs,ast) -> First $
case selectSmallestContaining (sp fs) ast of
Nothing -> Nothing
Just ast' -> Just ast'
where
sloc fs = mkRealSrcLoc fs sl sc
sp fs = mkRealSrcSpan (sloc fs) (sloc fs)
findEvidenceUse :: NodeIdentifiers a -> [Name]
findEvidenceUse ni = [n | (Right n, dets) <- xs, any isEvidenceUse (identInfo dets)]
where
xs = M.toList ni
data EvidenceInfo a
= EvidenceInfo
{ evidenceVar :: Name
, evidenceSpan :: RealSrcSpan
, evidenceType :: a
, evidenceDetails :: Maybe (EvVarSource, Scope, Maybe Span)
} deriving (Eq,Ord,Functor)
instance (Outputable a) => Outputable (EvidenceInfo a) where
ppr (EvidenceInfo name span typ dets) =
hang (ppr name <+> text "at" <+> ppr span O.<> text ", of type:" <+> ppr typ) 4 $
pdets $$ (pprDefinedAt name)
where
pdets = case dets of
Nothing -> text "is a usage of an external evidence variable"
Just (src,scp,spn) -> text "is an" <+> ppr (EvidenceVarBind src scp spn)
getEvidenceTreesAtPoint :: HieFile -> RefMap a -> (Int,Int) -> Tree.Forest (EvidenceInfo a)
getEvidenceTreesAtPoint hf refmap point =
[t | Just ast <- pure $ selectPoint hf point
, n <- findEvidenceUse (sourcedNodeIdents $ sourcedNodeInfo ast)
, Just t <- pure $ getEvidenceTree refmap n
]
getEvidenceTree :: RefMap a -> Name -> Maybe (Tree.Tree (EvidenceInfo a))
getEvidenceTree refmap var = go emptyNameSet var
where
go seen var
| var `elemNameSet` seen = Nothing
| otherwise = do
xs <- M.lookup (Right var) refmap
case find (any isEvidenceBind . identInfo . snd) xs of
Just (sp,dets) -> do
typ <- identType dets
(evdet,children) <- getFirst $ foldMap First $ do
det <- S.toList $ identInfo dets
case det of
EvidenceVarBind src@(EvLetBind (getEvBindDeps -> xs)) scp spn ->
pure $ Just ((src,scp,spn),mapMaybe (go $ extendNameSet seen var) xs)
EvidenceVarBind src scp spn -> pure $ Just ((src,scp,spn),[])
_ -> pure Nothing
pure $ Tree.Node (EvidenceInfo var sp typ (Just evdet)) children
Nothing -> getFirst $ foldMap First $ do
(sp,dets) <- xs
if (any isEvidenceUse $ identInfo dets)
then do
case identType dets of
Nothing -> pure Nothing
Just typ -> pure $ Just $ Tree.Node (EvidenceInfo var sp typ Nothing) []
else pure Nothing
hieTypeToIface :: HieTypeFix -> IfaceType
hieTypeToIface = foldType go
where
go (HTyVarTy n) = IfaceTyVar $ occNameFS $ getOccName n
go (HAppTy a b) = IfaceAppTy a (hieToIfaceArgs b)
go (HLitTy l) = IfaceLitTy l
go (HForAllTy ((n,k),af) t) = let b = (occNameFS $ getOccName n, k)
in IfaceForAllTy (Bndr (IfaceTvBndr b) af) t
go (HFunTy w a b) = IfaceFunTy VisArg w a b
go (HQualTy pred b) = IfaceFunTy InvisArg many_ty pred b
go (HCastTy a) = a
go HCoercionTy = IfaceTyVar "<coercion type>"
go (HTyConApp a xs) = IfaceTyConApp a (hieToIfaceArgs xs)
hieToIfaceArgs :: HieArgs IfaceType -> IfaceAppArgs
hieToIfaceArgs (HieArgs xs) = go' xs
where
go' [] = IA_Nil
go' ((True ,x):xs) = IA_Arg x Required $ go' xs
go' ((False,x):xs) = IA_Arg x Specified $ go' xs
data HieTypeState
= HTS
{ tyMap :: !(TypeMap TypeIndex)
, htyTable :: !(IM.IntMap HieTypeFlat)
, freshIndex :: !TypeIndex
}
initialHTS :: HieTypeState
initialHTS = HTS emptyTypeMap IM.empty 0
freshTypeIndex :: State HieTypeState TypeIndex
freshTypeIndex = do
index <- gets freshIndex
modify' $ \hts -> hts { freshIndex = index+1 }
return index
compressTypes
:: HieASTs Type
-> (HieASTs TypeIndex, A.Array TypeIndex HieTypeFlat)
compressTypes asts = (a, arr)
where
(a, (HTS _ m i)) = flip runState initialHTS $
for asts $ \typ -> do
i <- getTypeIndex typ
return i
arr = A.array (0,i1) (IM.toList m)
recoverFullType :: TypeIndex -> A.Array TypeIndex HieTypeFlat -> HieTypeFix
recoverFullType i m = go i
where
go i = Roll $ fmap go (m A.! i)
getTypeIndex :: Type -> State HieTypeState TypeIndex
getTypeIndex t
| otherwise = do
tm <- gets tyMap
case lookupTypeMap tm t of
Just i -> return i
Nothing -> do
ht <- go t
extendHTS t ht
where
extendHTS t ht = do
i <- freshTypeIndex
modify' $ \(HTS tm tt fi) ->
HTS (extendTypeMap tm t i) (IM.insert i ht tt) fi
return i
go (TyVarTy v) = return $ HTyVarTy $ varName v
go ty@(AppTy _ _) = do
let (head,args) = splitAppTys ty
visArgs = HieArgs $ resolveVisibility (typeKind head) args
ai <- getTypeIndex head
argsi <- mapM getTypeIndex visArgs
return $ HAppTy ai argsi
go (TyConApp f xs) = do
let visArgs = HieArgs $ resolveVisibility (tyConKind f) xs
is <- mapM getTypeIndex visArgs
return $ HTyConApp (toIfaceTyCon f) is
go (ForAllTy (Bndr v a) t) = do
k <- getTypeIndex (varType v)
i <- getTypeIndex t
return $ HForAllTy ((varName v,k),a) i
go (FunTy { ft_af = af, ft_mult = w, ft_arg = a, ft_res = b }) = do
ai <- getTypeIndex a
bi <- getTypeIndex b
wi <- getTypeIndex w
return $ case af of
InvisArg -> case w of Many -> HQualTy ai bi; _ -> error "Unexpected non-unrestricted predicate"
VisArg -> HFunTy wi ai bi
go (LitTy a) = return $ HLitTy $ toIfaceTyLit a
go (CastTy t _) = do
i <- getTypeIndex t
return $ HCastTy i
go (CoercionTy _) = return HCoercionTy
resolveTyVarScopes :: M.Map FastString (HieAST a) -> M.Map FastString (HieAST a)
resolveTyVarScopes asts = M.map go asts
where
go ast = resolveTyVarScopeLocal ast asts
resolveTyVarScopeLocal :: HieAST a -> M.Map FastString (HieAST a) -> HieAST a
resolveTyVarScopeLocal ast asts = go ast
where
resolveNameScope dets = dets{identInfo =
S.map resolveScope (identInfo dets)}
resolveScope (TyVarBind sc (UnresolvedScope names Nothing)) =
TyVarBind sc $ ResolvedScopes
[ LocalScope binding
| name <- names
, Just binding <- [getNameBinding name asts]
]
resolveScope (TyVarBind sc (UnresolvedScope names (Just sp))) =
TyVarBind sc $ ResolvedScopes
[ LocalScope binding
| name <- names
, Just binding <- [getNameBindingInClass name sp asts]
]
resolveScope scope = scope
go (Node info span children) = Node info' span $ map go children
where
info' = SourcedNodeInfo (updateNodeInfo <$> getSourcedNodeInfo info)
updateNodeInfo i = i { nodeIdentifiers = idents }
where
idents = M.map resolveNameScope $ nodeIdentifiers i
getNameBinding :: Name -> M.Map FastString (HieAST a) -> Maybe Span
getNameBinding n asts = do
(_,msp) <- getNameScopeAndBinding n asts
msp
getNameScope :: Name -> M.Map FastString (HieAST a) -> Maybe [Scope]
getNameScope n asts = do
(scopes,_) <- getNameScopeAndBinding n asts
return scopes
getNameBindingInClass
:: Name
-> Span
-> M.Map FastString (HieAST a)
-> Maybe Span
getNameBindingInClass n sp asts = do
ast <- M.lookup (srcSpanFile sp) asts
getFirst $ foldMap First $ do
child <- flattenAst ast
dets <- maybeToList
$ M.lookup (Right n) $ sourcedNodeIdents $ sourcedNodeInfo child
let binding = foldMap (First . getBindSiteFromContext) (identInfo dets)
return (getFirst binding)
getNameScopeAndBinding
:: Name
-> M.Map FastString (HieAST a)
-> Maybe ([Scope], Maybe Span)
getNameScopeAndBinding n asts = case nameSrcSpan n of
RealSrcSpan sp _ -> do
ast <- M.lookup (srcSpanFile sp) asts
defNode <- selectLargestContainedBy sp ast
getFirst $ foldMap First $ do
node <- flattenAst defNode
dets <- maybeToList
$ M.lookup (Right n) $ sourcedNodeIdents $ sourcedNodeInfo node
scopes <- maybeToList $ foldMap getScopeFromContext (identInfo dets)
let binding = foldMap (First . getBindSiteFromContext) (identInfo dets)
return $ Just (scopes, getFirst binding)
_ -> Nothing
getScopeFromContext :: ContextInfo -> Maybe [Scope]
getScopeFromContext (ValBind _ sc _) = Just [sc]
getScopeFromContext (PatternBind a b _) = Just [a, b]
getScopeFromContext (ClassTyDecl _) = Just [ModuleScope]
getScopeFromContext (Decl _ _) = Just [ModuleScope]
getScopeFromContext (TyVarBind a (ResolvedScopes xs)) = Just $ a:xs
getScopeFromContext (TyVarBind a _) = Just [a]
getScopeFromContext (EvidenceVarBind _ a _) = Just [a]
getScopeFromContext _ = Nothing
getBindSiteFromContext :: ContextInfo -> Maybe Span
getBindSiteFromContext (ValBind _ _ sp) = sp
getBindSiteFromContext (PatternBind _ _ sp) = sp
getBindSiteFromContext _ = Nothing
flattenAst :: HieAST a -> [HieAST a]
flattenAst n =
n : concatMap flattenAst (nodeChildren n)
smallestContainingSatisfying
:: Span
-> (HieAST a -> Bool)
-> HieAST a
-> Maybe (HieAST a)
smallestContainingSatisfying sp cond node
| nodeSpan node `containsSpan` sp = getFirst $ mconcat
[ foldMap (First . smallestContainingSatisfying sp cond) $
nodeChildren node
, First $ if cond node then Just node else Nothing
]
| sp `containsSpan` nodeSpan node = Nothing
| otherwise = Nothing
selectLargestContainedBy :: Span -> HieAST a -> Maybe (HieAST a)
selectLargestContainedBy sp node
| sp `containsSpan` nodeSpan node = Just node
| nodeSpan node `containsSpan` sp =
getFirst $ foldMap (First . selectLargestContainedBy sp) $
nodeChildren node
| otherwise = Nothing
selectSmallestContaining :: Span -> HieAST a -> Maybe (HieAST a)
selectSmallestContaining sp node
| nodeSpan node `containsSpan` sp = getFirst $ mconcat
[ foldMap (First . selectSmallestContaining sp) $ nodeChildren node
, First (Just node)
]
| sp `containsSpan` nodeSpan node = Nothing
| otherwise = Nothing
definedInAsts :: M.Map FastString (HieAST a) -> Name -> Bool
definedInAsts asts n = case nameSrcSpan n of
RealSrcSpan sp _ -> srcSpanFile sp `elem` M.keys asts
_ -> False
getEvidenceBindDeps :: ContextInfo -> [Name]
getEvidenceBindDeps (EvidenceVarBind (EvLetBind xs) _ _) =
getEvBindDeps xs
getEvidenceBindDeps _ = []
isEvidenceBind :: ContextInfo -> Bool
isEvidenceBind EvidenceVarBind{} = True
isEvidenceBind _ = False
isEvidenceContext :: ContextInfo -> Bool
isEvidenceContext EvidenceVarUse = True
isEvidenceContext EvidenceVarBind{} = True
isEvidenceContext _ = False
isEvidenceUse :: ContextInfo -> Bool
isEvidenceUse EvidenceVarUse = True
isEvidenceUse _ = False
isOccurrence :: ContextInfo -> Bool
isOccurrence Use = True
isOccurrence EvidenceVarUse = True
isOccurrence _ = False
scopeContainsSpan :: Scope -> Span -> Bool
scopeContainsSpan NoScope _ = False
scopeContainsSpan ModuleScope _ = True
scopeContainsSpan (LocalScope a) b = a `containsSpan` b
combineAst :: HieAST Type -> HieAST Type -> HieAST Type
combineAst a@(Node aInf aSpn xs) b@(Node bInf bSpn ys)
| aSpn == bSpn = Node (aInf `combineSourcedNodeInfo` bInf) aSpn (mergeAsts xs ys)
| aSpn `containsSpan` bSpn = combineAst b a
combineAst a (Node xs span children) = Node xs span (insertAst a children)
insertAst :: HieAST Type -> [HieAST Type] -> [HieAST Type]
insertAst x = mergeAsts [x]
nodeInfo :: HieAST Type -> NodeInfo Type
nodeInfo = foldl' combineNodeInfo emptyNodeInfo . getSourcedNodeInfo . sourcedNodeInfo
emptyNodeInfo :: NodeInfo a
emptyNodeInfo = NodeInfo S.empty [] M.empty
sourcedNodeIdents :: SourcedNodeInfo a -> NodeIdentifiers a
sourcedNodeIdents = M.unionsWith (<>) . fmap nodeIdentifiers . getSourcedNodeInfo
combineSourcedNodeInfo :: SourcedNodeInfo Type -> SourcedNodeInfo Type -> SourcedNodeInfo Type
combineSourcedNodeInfo = coerce $ M.unionWith combineNodeInfo
combineNodeInfo :: NodeInfo Type -> NodeInfo Type -> NodeInfo Type
(NodeInfo as ai ad) `combineNodeInfo` (NodeInfo bs bi bd) =
NodeInfo (S.union as bs) (mergeSorted ai bi) (M.unionWith (<>) ad bd)
where
mergeSorted :: [Type] -> [Type] -> [Type]
mergeSorted la@(a:as) lb@(b:bs) = case nonDetCmpType a b of
LT -> a : mergeSorted as lb
EQ -> a : mergeSorted as bs
GT -> b : mergeSorted la bs
mergeSorted as [] = as
mergeSorted [] bs = bs
mergeAsts :: [HieAST Type] -> [HieAST Type] -> [HieAST Type]
mergeAsts xs [] = xs
mergeAsts [] ys = ys
mergeAsts xs@(a:as) ys@(b:bs)
| span_a `containsSpan` span_b = mergeAsts (combineAst a b : as) bs
| span_b `containsSpan` span_a = mergeAsts as (combineAst a b : bs)
| span_a `rightOf` span_b = b : mergeAsts xs bs
| span_a `leftOf` span_b = a : mergeAsts as ys
| span_a `startsRightOf` span_b = b : mergeAsts as ys
| otherwise = a : mergeAsts as ys
where
span_a = nodeSpan a
span_b = nodeSpan b
rightOf :: Span -> Span -> Bool
rightOf s1 s2
= (srcSpanStartLine s1, srcSpanStartCol s1)
>= (srcSpanEndLine s2, srcSpanEndCol s2)
&& (srcSpanFile s1 == srcSpanFile s2)
leftOf :: Span -> Span -> Bool
leftOf s1 s2
= (srcSpanEndLine s1, srcSpanEndCol s1)
<= (srcSpanStartLine s2, srcSpanStartCol s2)
&& (srcSpanFile s1 == srcSpanFile s2)
startsRightOf :: Span -> Span -> Bool
startsRightOf s1 s2
= (srcSpanStartLine s1, srcSpanStartCol s1)
>= (srcSpanStartLine s2, srcSpanStartCol s2)
mergeSortAsts :: [HieAST Type] -> [HieAST Type]
mergeSortAsts = go . map pure
where
go [] = []
go [xs] = xs
go xss = go (mergePairs xss)
mergePairs [] = []
mergePairs [xs] = [xs]
mergePairs (xs:ys:xss) = mergeAsts xs ys : mergePairs xss
simpleNodeInfo :: FastString -> FastString -> NodeInfo a
simpleNodeInfo cons typ = NodeInfo (S.singleton (cons, typ)) [] M.empty
locOnly :: Monad m => SrcSpan -> ReaderT NodeOrigin m [HieAST a]
locOnly (RealSrcSpan span _) = do
org <- ask
let e = mkSourcedNodeInfo org $ emptyNodeInfo
pure [Node e span []]
locOnly _ = pure []
mkScope :: SrcSpan -> Scope
mkScope (RealSrcSpan sp _) = LocalScope sp
mkScope _ = NoScope
mkLScope :: Located a -> Scope
mkLScope = mkScope . getLoc
combineScopes :: Scope -> Scope -> Scope
combineScopes ModuleScope _ = ModuleScope
combineScopes _ ModuleScope = ModuleScope
combineScopes NoScope x = x
combineScopes x NoScope = x
combineScopes (LocalScope a) (LocalScope b) =
mkScope $ combineSrcSpans (RealSrcSpan a Nothing) (RealSrcSpan b Nothing)
mkSourcedNodeInfo :: NodeOrigin -> NodeInfo a -> SourcedNodeInfo a
mkSourcedNodeInfo org ni = SourcedNodeInfo $ M.singleton org ni
makeNode
:: (Monad m, Data a)
=> a
-> SrcSpan
-> ReaderT NodeOrigin m [HieAST b]
makeNode x spn = do
org <- ask
pure $ case spn of
RealSrcSpan span _ -> [Node (mkSourcedNodeInfo org $ simpleNodeInfo cons typ) span []]
_ -> []
where
cons = mkFastString . show . toConstr $ x
typ = mkFastString . show . typeRepTyCon . typeOf $ x
makeTypeNode
:: (Monad m, Data a)
=> a
-> SrcSpan
-> Type
-> ReaderT NodeOrigin m [HieAST Type]
makeTypeNode x spn etyp = do
org <- ask
pure $ case spn of
RealSrcSpan span _ ->
[Node (mkSourcedNodeInfo org $ NodeInfo (S.singleton (cons,typ)) [etyp] M.empty) span []]
_ -> []
where
cons = mkFastString . show . toConstr $ x
typ = mkFastString . show . typeRepTyCon . typeOf $ x