{-# 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