From 2d489a92b5ff9bbf92d3f4a5177a8dcf5aeb2e99 Mon Sep 17 00:00:00 2001 From: Simon Jakobi Date: Fri, 3 Apr 2020 21:27:33 +0200 Subject: Join the binding traversals --- compiler/GHC/Stg/DepAnal.hs | 206 +++++++++++++++++++++++++++++++------------- 1 file changed, 145 insertions(+), 61 deletions(-) diff --git a/compiler/GHC/Stg/DepAnal.hs b/compiler/GHC/Stg/DepAnal.hs index 5947006e0e..d811277f0f 100644 --- a/compiler/GHC/Stg/DepAnal.hs +++ b/compiler/GHC/Stg/DepAnal.hs @@ -4,7 +4,7 @@ module GHC.Stg.DepAnal (depSortStgPgm) where import GhcPrelude -import GHC.Stg.FVs +import GHC.Core ( Tickish(Breakpoint) ) import GHC.Stg.Syntax import GHC.Types.Id import GHC.Types.Name (Name, nameIsLocalOrFrom) @@ -13,6 +13,9 @@ import Outputable import GHC.Types.Unique.Set (nonDetEltsUniqSet) import GHC.Types.Var.Set import GHC.Types.Module (Module) +import Util + +import Data.Maybe ( catMaybes ) import Data.Graph (SCC (..)) @@ -25,7 +28,23 @@ type BVs = VarSet -- | Set of free variables type FVs = VarSet --- | Dependency analysis on STG terms. +newtype Env + = Env + { locals :: IdSet + } + +emptyEnv :: Env +emptyEnv = Env emptyVarSet + +addLocals :: [Id] -> Env -> Env +addLocals bndrs env + = env { locals = extendVarSetList (locals env) bndrs } + +-- | This makes sure that only local, non-global free vars make it into the set. +mkFreeVarSet :: Env -> [Id] -> DIdSet +mkFreeVarSet env = mkDVarSet . filter (`elemVarSet` locals env) + +-- | Dependency analysis and free variable annotations on STG terms. -- -- Dependencies of a binding are just free variables in the binding. This -- includes imported ids and ids in the current module. For recursive groups we @@ -44,25 +63,130 @@ annTopBindingsDeps this_mod bs = map top_bind bs (StgTopStringLit id bs, emptyVarSet) top_bind (StgTopLifted bs) = - (StgTopLifted (annBindingFreeVars bs), binding emptyVarSet bs) - - binding :: BVs -> StgBinding -> FVs - binding bounds (StgNonRec _ r) = - rhs bounds r - binding bounds (StgRec bndrs) = - unionVarSets $ - map (bind_non_rec (extendVarSetList bounds (map fst bndrs))) bndrs - - bind_non_rec :: BVs -> (Id, StgRhs) -> FVs - bind_non_rec bounds (_, r) = - rhs bounds r - - rhs :: BVs -> StgRhs -> FVs - rhs bounds (StgRhsClosure _ _ _ as e) = - expr (extendVarSetList bounds as) e - - rhs bounds (StgRhsCon _ _ as) = - args bounds as + (StgTopLifted bs', fvs) + where + (bs', _dIdSet, fvs) = binding emptyEnv emptyDVarSet emptyVarSet bs + + binding :: Env -> DIdSet -> BVs -> StgBinding -> (CgStgBinding, DIdSet, FVs) + binding env body_fv bounds (StgNonRec bndr r) = + (StgNonRec bndr r', fvs, da_fvs) + where + -- See Note [Tracking local binders] + (r', rhs_fvs, da_fvs) = rhs env bounds r + fvs = delDVarSet body_fv bndr `unionDVarSet` rhs_fvs + binding env body_fv bounds (StgRec bindings) = + ( StgRec (zip bndrs rhss') + , delDVarSetList (unionDVarSets (body_fv:rhs_fvss)) bndrs + , unionVarSets da_fvss + ) + where + (bndrs, rhss) = unzip bindings + bounds' = extendVarSetList bounds bndrs + (rhss', rhs_fvss, da_fvss) = mapAndUnzip3 (rhs env bounds') rhss + + rhs :: Env -> BVs -> StgRhs -> (CgStgRhs, DIdSet, FVs) + rhs env bounds (StgRhsClosure _ ccs uf bndrs body) = + ( StgRhsClosure fvs ccs uf bndrs body' + , fvs + , da_fvs + ) + where + (body', body_fvs, da_fvs) = expr (addLocals bndrs env) (extendVarSetList bounds bndrs) body + fvs = delDVarSetList body_fvs bndrs + rhs env bounds (StgRhsCon ccs dc as) = + ( StgRhsCon ccs dc as + , fvs + , da_fvs + ) + where + (fvs, da_fvs) = args env bounds as + + expr :: Env -> BVs -> StgExpr -> (CgStgExpr, DIdSet, FVs) + expr env = go + where + go bounds (StgApp occ as) = + ( StgApp occ as + , unionDVarSet fvs (mkFreeVarSet env [occ]) + , var bounds occ `unionVarSet` da_fvs + ) + where + (fvs, da_fvs) = args env bounds as + go _ (StgLit lit) = + (StgLit lit, emptyDVarSet, emptyVarSet) + go bounds (StgConApp dc as tys) = + (StgConApp dc as tys, fvs, da_fvs) + where + (fvs, da_fvs) = args env bounds as + go bounds (StgOpApp op as ty) = + (StgOpApp op as ty, fvs, da_fvs) + where + (fvs, da_fvs) = args env bounds as + go _ lam@StgLam{} = + pprPanic "annTopBindingsDeps" (text "Found lambda:" $$ ppr lam) + go bounds (StgCase scrut bndr ty as) = + ( StgCase scrut' bndr ty alts' + , delDVarSet (unionDVarSet scrut_fvs alt_fvs) bndr + , scrut_da_fvs `unionVarSet` alt_da_fvs + ) + where + (scrut', scrut_fvs, scrut_da_fvs) = go bounds scrut + -- See Note [Tracking local binders] + (alts', alt_fvs, alt_da_fvs) = + alts (addLocals [bndr] env) (extendVarSet bounds bndr) as + go bounds (StgLet ext bind body) = + go_bind bounds (StgLet ext) bind body + go bounds (StgLetNoEscape ext bind body) = + go_bind bounds (StgLetNoEscape ext) bind body + go bounds (StgTick tick e) = + (StgTick tick e', fvs', da_fvs) + where + (e', fvs, da_fvs) = go bounds e + fvs' = unionDVarSet (tickish tick) fvs + tickish (Breakpoint _ ids) = mkDVarSet ids + tickish _ = emptyDVarSet + + go_bind bounds dc bind body = + ( dc bind' body' + , fvs + , da_bind_fvs `unionVarSet` da_body_fvs + ) + where + -- See Note [Tracking local binders] + binders = bindersOf bind + env' = addLocals binders env + (body', body_fvs, da_body_fvs) = + expr env' (extendVarSetList bounds binders) body + (bind', fvs, da_bind_fvs) = binding env' body_fvs bounds bind + + alts :: Env -> BVs -> [StgAlt] -> ([CgStgAlt], DIdSet, FVs) + alts env bounds as = + ( as' + , unionDVarSets alt_fvss + , unionVarSets alt_da_fvss + ) + where + (as', alt_fvss, alt_da_fvss) = mapAndUnzip3 (alt env bounds) as + + alt :: Env -> BVs -> StgAlt -> (CgStgAlt, DIdSet, FVs) + alt env bounds (con, bndrs, e) = + ( (con, bndrs, e') + , delDVarSetList rhs_fvs bndrs + , da_fvs + ) + where + (e', rhs_fvs, da_fvs) = expr (addLocals bndrs env) (extendVarSetList bounds bndrs) e + + args :: Env -> BVs -> [StgArg] -> (DIdSet, FVs) + args env bounds as = + ( mkFreeVarSet env (catMaybes mIds) + , unionVarSets da_fvss + ) + where + (mIds, da_fvss) = mapAndUnzip (arg bounds) as + + arg :: BVs -> StgArg -> (Maybe Id, FVs) + arg bounds (StgVarArg v) = (Just v, var bounds v) + arg _ StgLitArg{} = (Nothing, emptyVarSet) var :: BVs -> Var -> FVs var bounds v @@ -72,46 +196,6 @@ annTopBindingsDeps this_mod bs = map top_bind bs | otherwise = emptyVarSet - arg :: BVs -> StgArg -> FVs - arg bounds (StgVarArg v) = var bounds v - arg _ StgLitArg{} = emptyVarSet - - args :: BVs -> [StgArg] -> FVs - args bounds as = unionVarSets (map (arg bounds) as) - - expr :: BVs -> StgExpr -> FVs - expr bounds (StgApp f as) = - var bounds f `unionVarSet` args bounds as - - expr _ StgLit{} = - emptyVarSet - - expr bounds (StgConApp _ as _) = - args bounds as - expr bounds (StgOpApp _ as _) = - args bounds as - expr _ lam@StgLam{} = - pprPanic "annTopBindingsDeps" (text "Found lambda:" $$ ppr lam) - expr bounds (StgCase scrut scrut_bndr _ as) = - expr bounds scrut `unionVarSet` - alts (extendVarSet bounds scrut_bndr) as - expr bounds (StgLet _ bs e) = - binding bounds bs `unionVarSet` - expr (extendVarSetList bounds (bindersOf bs)) e - expr bounds (StgLetNoEscape _ bs e) = - binding bounds bs `unionVarSet` - expr (extendVarSetList bounds (bindersOf bs)) e - - expr bounds (StgTick _ e) = - expr bounds e - - alts :: BVs -> [StgAlt] -> FVs - alts bounds = unionVarSets . map (alt bounds) - - alt :: BVs -> StgAlt -> FVs - alt bounds (_, bndrs, e) = - expr (extendVarSetList bounds bndrs) e - -------------------------------------------------------------------------------- -- * Dependency sorting -- cgit v1.2.1