{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
module GHC.Util.FreeVars (
vars', varss', pvars',
Vars' (..), FreeVars'(..) , AllVars' (..)
) where
import RdrName
import GHC.Hs.Types
import OccName
import Name
import GHC.Hs
import SrcLoc
import Bag (bagToList)
import Data.Generics.Uniplate.Data ()
import Data.Generics.Uniplate.Operations
import Data.Monoid
import Data.Semigroup
import Data.List.Extra
import Data.Set (Set)
import qualified Data.Set as Set
import Prelude
( ^+ ) :: Set OccName -> Set OccName -> Set OccName
( ^+ ) = Set.union
( ^- ) :: Set OccName -> Set OccName -> Set OccName
( ^- ) = Set.difference
data Vars' = Vars'{bound' :: Set OccName, free' :: Set OccName}
instance Show Vars' where
show (Vars' bs fs) = "bound : " ++
show (map occNameString (Set.toList bs)) ++
", free : " ++ show (map occNameString (Set.toList fs))
instance Semigroup Vars' where
Vars' x1 x2 <> Vars' y1 y2 = Vars' (x1 ^+ y1) (x2 ^+ y2)
instance Monoid Vars' where
mempty = Vars' Set.empty Set.empty
mconcat vs = Vars' (Set.unions $ map bound' vs) (Set.unions $ map free' vs)
class AllVars' a where
allVars' :: a -> Vars'
class FreeVars' a where
freeVars' :: a -> Set OccName
instance AllVars' Vars' where allVars' = id
instance FreeVars' (Set OccName) where freeVars' = id
instance (AllVars' a) => AllVars' [a] where allVars' = mconcatMap allVars'
instance (FreeVars' a) => FreeVars' [a] where freeVars' = Set.unions . map freeVars'
freeVars_' :: (FreeVars' a) => a -> Vars'
freeVars_' = Vars' Set.empty . freeVars'
inFree' :: (AllVars' a, FreeVars' b) => a -> b -> Set OccName
inFree' a b = free' aa ^+ (freeVars' b ^- bound' aa)
where aa = allVars' a
inVars' :: (AllVars' a, AllVars' b) => a -> b -> Vars'
inVars' a b =
Vars' (bound' aa ^+ bound' bb) (free' aa ^+ (free' bb ^- bound' aa))
where aa = allVars' a
bb = allVars' b
unqualNames' :: Located RdrName -> [OccName]
unqualNames' (L _ (Unqual x)) = [x]
unqualNames' (L _ (Exact x)) = [nameOccName x]
unqualNames' _ = []
instance FreeVars' (LHsExpr GhcPs) where
freeVars' (L _ (HsVar _ x)) = Set.fromList $ unqualNames' x
freeVars' (L _ (HsUnboundVar _ x)) = Set.fromList [unboundVarOcc x]
freeVars' (L _ (HsLam _ mg)) = free' (allVars' mg)
freeVars' (L _ (HsLamCase _ MG{mg_alts=(L _ ms)})) = free' (allVars' ms)
freeVars' (L _ (HsCase _ of_ MG{mg_alts=(L _ ms)})) = freeVars' of_ ^+ free' (allVars' ms)
freeVars' (L _ (HsLet _ binds e)) = inFree' binds e
freeVars' (L _ (HsDo _ ctxt (L _ stmts))) = free' (allVars' stmts)
freeVars' (L _ (RecordCon _ _ (HsRecFields flds _))) = Set.unions $ map freeVars' flds
freeVars' (L _ (RecordUpd _ e flds)) = Set.unions $ freeVars' e : map freeVars' flds
freeVars' (L _ (HsMultiIf _ grhss)) = free' (allVars' grhss)
freeVars' (L _ HsConLikeOut{}) = mempty
freeVars' (L _ HsRecFld{}) = mempty
freeVars' (L _ HsOverLabel{}) = mempty
freeVars' (L _ HsIPVar{}) = mempty
freeVars' (L _ HsOverLit{}) = mempty
freeVars' (L _ HsLit{}) = mempty
freeVars' (L _ HsRnBracketOut{}) = mempty
freeVars' (L _ HsTcBracketOut{}) = mempty
freeVars' (L _ HsWrap{}) = mempty
freeVars' e = freeVars' $ children e
instance FreeVars' (LHsTupArg GhcPs) where
freeVars' (L _ (Present _ args)) = freeVars' args
freeVars' _ = mempty
instance FreeVars' (LHsRecField GhcPs (LHsExpr GhcPs)) where
freeVars' o@(L _ (HsRecField x _ True)) = Set.singleton $ occName $ unLoc $ rdrNameFieldOcc $ unLoc x
freeVars' o@(L _ (HsRecField _ x _)) = freeVars' x
instance FreeVars' (LHsRecUpdField GhcPs) where
freeVars' (L _ (HsRecField _ x _)) = freeVars' x
instance AllVars' (Located (Pat GhcPs)) where
allVars' (L _ (VarPat _ (L _ x))) = Vars' (Set.singleton $ rdrNameOcc x) Set.empty
allVars' (L _ (AsPat _ n x)) = allVars' (noLoc $ VarPat noExtField n :: LPat GhcPs) <> allVars' x
allVars' (L _ (ConPatIn _ (RecCon (HsRecFields flds _)))) = allVars' flds
allVars' (L _ (NPlusKPat _ n _ _ _ _)) = allVars' (noLoc $ VarPat noExtField n :: LPat GhcPs)
allVars' (L _ (ViewPat _ e p)) = freeVars_' e <> allVars' p
allVars' (L _ WildPat{}) = mempty
allVars' (L _ ConPatOut{}) = mempty
allVars' (L _ LitPat{}) = mempty
allVars' (L _ NPat{}) = mempty
allVars' p = allVars' $ children p
instance AllVars' (LHsRecField GhcPs (Located (Pat GhcPs))) where
allVars' (L _ (HsRecField _ x _)) = allVars' x
instance AllVars' (LStmt GhcPs (LHsExpr GhcPs)) where
allVars' (L _ (LastStmt _ expr _ _)) = freeVars_' expr
allVars' (L _ (BindStmt _ pat expr _ _)) = allVars' pat <> freeVars_' expr
allVars' (L _ (BodyStmt _ expr _ _)) = freeVars_' expr
allVars' (L _ (LetStmt _ binds)) = allVars' binds
allVars' (L _ (TransStmt _ _ stmts _ using by _ _ fmap_)) = allVars' stmts <> freeVars_' using <> maybe mempty freeVars_' by <> freeVars_' (noLoc fmap_ :: Located (HsExpr GhcPs))
allVars' (L _ (RecStmt _ stmts _ _ _ _ _)) = allVars' stmts
allVars' (L _ ApplicativeStmt{}) = mempty
allVars' (L _ ParStmt{}) = mempty
allVars' _ = mempty
instance AllVars' (LHsLocalBinds GhcPs) where
allVars' (L _ (HsValBinds _ (ValBinds _ binds _))) = allVars' (bagToList binds)
allVars' (L _ (HsIPBinds _ (IPBinds _ binds))) = allVars' binds
allVars' (L _ EmptyLocalBinds{}) = mempty
allVars' _ = mempty
instance AllVars' (LIPBind GhcPs) where
allVars' (L _ (IPBind _ _ e)) = freeVars_' e
allVars' _ = mempty
instance AllVars' (LHsBind GhcPs) where
allVars' (L _ FunBind{fun_id=n, fun_matches=MG{mg_alts=(L _ ms)}}) = allVars' (noLoc $ VarPat noExtField n :: LPat GhcPs) <> allVars' ms
allVars' (L _ PatBind{pat_lhs=n, pat_rhs=grhss}) = allVars' n <> allVars' grhss
allVars' (L _ (PatSynBind _ PSB{})) = mempty
allVars' (L _ VarBind{}) = mempty
allVars' (L _ AbsBinds{}) = mempty
allVars' _ = mempty
instance AllVars' (MatchGroup GhcPs (LHsExpr GhcPs)) where
allVars' (MG _ _alts@(L _ alts) _) = inVars' (foldMap (allVars' . m_pats) ms) (allVars' (map m_grhss ms))
where ms = map unLoc alts
allVars' _ = mempty
instance AllVars' (LMatch GhcPs (LHsExpr GhcPs)) where
allVars' (L _ (Match _ FunRhs {mc_fun=name} pats grhss)) = allVars' (noLoc $ VarPat noExtField name :: LPat GhcPs) <> allVars' pats <> allVars' grhss
allVars' (L _ (Match _ (StmtCtxt ctxt) pats grhss)) = allVars' ctxt <> allVars' pats <> allVars' grhss
allVars' (L _ (Match _ _ pats grhss)) = inVars' (allVars' pats) (allVars' grhss)
allVars' _ = mempty
instance AllVars' (HsStmtContext RdrName) where
allVars' (PatGuard FunRhs{mc_fun=n}) = allVars' (noLoc $ VarPat noExtField n :: LPat GhcPs)
allVars' ParStmtCtxt{} = mempty
allVars' TransStmtCtxt{} = mempty
allVars' _ = mempty
instance AllVars' (GRHSs GhcPs (LHsExpr GhcPs)) where
allVars' (GRHSs _ grhss binds) = inVars' binds (mconcatMap allVars' grhss)
allVars' _ = mempty
instance AllVars' (LGRHS GhcPs (LHsExpr GhcPs)) where
allVars' (L _ (GRHS _ guards expr)) = Vars' (bound' gs) (free' gs ^+ (freeVars' expr ^- bound' gs)) where gs = allVars' guards
allVars' _ = mempty
instance AllVars' (LHsDecl GhcPs) where
allVars' (L l (ValD _ bind)) = allVars' (L l bind :: LHsBind GhcPs)
allVars' _ = mempty
vars' :: FreeVars' a => a -> [String]
vars' = Set.toList . Set.map occNameString . freeVars'
varss' :: AllVars' a => a -> [String]
varss' = Set.toList . Set.map occNameString . free' . allVars'
pvars' :: AllVars' a => a -> [String]
pvars' = Set.toList . Set.map occNameString . bound' . allVars'