module GHC.Stg.FVs (
depSortWithAnnotStgPgm,
annBindingFreeVars
) where
import GHC.Prelude hiding (mod)
import GHC.Stg.Syntax
import GHC.Stg.Utils (bindersOf)
import GHC.Types.Id
import GHC.Types.Name (Name, nameIsLocalOrFrom)
import GHC.Types.Tickish ( GenTickish(Breakpoint) )
import GHC.Types.Unique.Set (nonDetEltsUniqSet)
import GHC.Types.Var.Set
import GHC.Unit.Module (Module)
import GHC.Utils.Misc
import Data.Graph (SCC (..))
import GHC.Data.Graph.Directed( Node(..), stronglyConnCompFromEdgedVerticesUniq )
depSortWithAnnotStgPgm :: Module -> [StgTopBinding] -> [CgStgTopBinding]
depSortWithAnnotStgPgm this_mod binds
=
lit_binds ++ map from_scc sccs
where
lit_binds :: [CgStgTopBinding]
pairs :: [(Id, StgRhs)]
(lit_binds, pairs) = flattenTopStgBindings binds
nodes :: [Node Name (Id, CgStgRhs)]
nodes = map (annotateTopPair env0) pairs
env0 = Env { locals = emptyVarSet, mod = this_mod }
sccs :: [SCC (Id,CgStgRhs)]
sccs = stronglyConnCompFromEdgedVerticesUniq nodes
from_scc (CyclicSCC pairs) = StgTopLifted (StgRec pairs)
from_scc (AcyclicSCC (bndr,rhs)) = StgTopLifted (StgNonRec bndr rhs)
flattenTopStgBindings :: [StgTopBinding] -> ([CgStgTopBinding], [(Id,StgRhs)])
flattenTopStgBindings binds
= go [] [] binds
where
go lits pairs [] = (lits, pairs)
go lits pairs (bind:binds)
= case bind of
StgTopStringLit bndr rhs -> go (StgTopStringLit bndr rhs:lits) pairs binds
StgTopLifted stg_bind -> go lits (flatten_one stg_bind ++ pairs) binds
flatten_one (StgNonRec b r) = [(b,r)]
flatten_one (StgRec pairs) = pairs
annotateTopPair :: Env -> (Id, StgRhs) -> Node Name (Id, CgStgRhs)
annotateTopPair env0 (bndr, rhs)
= DigraphNode { node_key = idName bndr
, node_payload = (bndr, rhs')
, node_dependencies = map idName (nonDetEltsUniqSet top_fvs) }
where
(rhs', top_fvs, _) = rhsFVs env0 rhs
data Env
= Env
{
locals :: IdSet
, mod :: Module
}
addLocals :: [Id] -> Env -> Env
addLocals bndrs env
= env { locals = extendVarSetList (locals env) bndrs }
type TopFVs = IdSet
type LocalFVs = DIdSet
annBindingFreeVars :: Module -> StgBinding -> CgStgBinding
annBindingFreeVars this_mod = fstOf3 . bindingFVs (Env emptyVarSet this_mod) emptyDVarSet
bindingFVs :: Env -> LocalFVs -> StgBinding -> (CgStgBinding, TopFVs, LocalFVs)
bindingFVs env body_fv b =
case b of
StgNonRec bndr r -> (StgNonRec bndr r', fvs, lcl_fvs)
where
(r', fvs, rhs_lcl_fvs) = rhsFVs env r
lcl_fvs = delDVarSet body_fv bndr `unionDVarSet` rhs_lcl_fvs
StgRec pairs -> (StgRec pairs', fvs, lcl_fvss)
where
bndrs = map fst pairs
env' = addLocals bndrs env
(rhss, rhs_fvss, rhs_lcl_fvss) = mapAndUnzip3 (rhsFVs env' . snd) pairs
fvs = unionVarSets rhs_fvss
pairs' = zip bndrs rhss
lcl_fvss = delDVarSetList (unionDVarSets (body_fv:rhs_lcl_fvss)) bndrs
varFVs :: Env -> Id -> (TopFVs, LocalFVs) -> (TopFVs, LocalFVs)
varFVs env v (top_fvs, lcl_fvs)
| v `elemVarSet` locals env
= (top_fvs, lcl_fvs `extendDVarSet` v)
| nameIsLocalOrFrom (mod env) (idName v)
= (top_fvs `extendVarSet` v, lcl_fvs)
| otherwise
= (top_fvs, lcl_fvs)
exprFVs :: Env -> StgExpr -> (CgStgExpr, TopFVs, LocalFVs)
exprFVs env = go
where
go (StgApp f as)
| (top_fvs, lcl_fvs) <- varFVs env f (argsFVs env as)
= (StgApp f as, top_fvs, lcl_fvs)
go (StgLit lit) = (StgLit lit, emptyVarSet, emptyDVarSet)
go (StgConApp dc n as tys)
| (top_fvs, lcl_fvs) <- argsFVs env as
= (StgConApp dc n as tys, top_fvs, lcl_fvs)
go (StgOpApp op as ty)
| (top_fvs, lcl_fvs) <- argsFVs env as
= (StgOpApp op as ty, top_fvs, lcl_fvs)
go (StgCase scrut bndr ty alts)
| (scrut',scrut_top_fvs,scrut_lcl_fvs) <- exprFVs env scrut
, (alts',alts_top_fvss,alts_lcl_fvss)
<- mapAndUnzip3 (altFVs (addLocals [bndr] env)) alts
, let top_fvs = scrut_top_fvs `unionVarSet` unionVarSets alts_top_fvss
alts_lcl_fvs = unionDVarSets alts_lcl_fvss
lcl_fvs = delDVarSet (unionDVarSet scrut_lcl_fvs alts_lcl_fvs) bndr
= (StgCase scrut' bndr ty alts', top_fvs,lcl_fvs)
go (StgLet ext bind body) = go_bind (StgLet ext) bind body
go (StgLetNoEscape ext bind body) = go_bind (StgLetNoEscape ext) bind body
go (StgTick tick e)
| (e', top_fvs, lcl_fvs) <- exprFVs env e
, let lcl_fvs' = unionDVarSet (tickish tick) lcl_fvs
= (StgTick tick e', top_fvs, lcl_fvs')
where
tickish (Breakpoint _ _ ids) = mkDVarSet ids
tickish _ = emptyDVarSet
go_bind dc bind body = (dc bind' body', top_fvs, lcl_fvs)
where
env' = addLocals (bindersOf bind) env
(body', body_top_fvs, body_lcl_fvs) = exprFVs env' body
(bind', bind_top_fvs, lcl_fvs) = bindingFVs env' body_lcl_fvs bind
top_fvs = bind_top_fvs `unionVarSet` body_top_fvs
rhsFVs :: Env -> StgRhs -> (CgStgRhs, TopFVs, LocalFVs)
rhsFVs env (StgRhsClosure _ ccs uf bs body)
| (body', top_fvs, lcl_fvs) <- exprFVs (addLocals bs env) body
, let lcl_fvs' = delDVarSetList lcl_fvs bs
= (StgRhsClosure lcl_fvs' ccs uf bs body', top_fvs, lcl_fvs')
rhsFVs env (StgRhsCon ccs dc mu ts bs)
| (top_fvs, lcl_fvs) <- argsFVs env bs
= (StgRhsCon ccs dc mu ts bs, top_fvs, lcl_fvs)
argsFVs :: Env -> [StgArg] -> (TopFVs, LocalFVs)
argsFVs env = foldl' f (emptyVarSet, emptyDVarSet)
where
f (fvs,ids) StgLitArg{} = (fvs, ids)
f (fvs,ids) (StgVarArg v) = varFVs env v (fvs, ids)
altFVs :: Env -> StgAlt -> (CgStgAlt, TopFVs, LocalFVs)
altFVs env GenStgAlt{alt_con=con, alt_bndrs=bndrs, alt_rhs=e}
| (e', top_fvs, lcl_fvs) <- exprFVs (addLocals bndrs env) e
, let lcl_fvs' = delDVarSetList lcl_fvs bndrs
, let newAlt = GenStgAlt{alt_con=con, alt_bndrs=bndrs, alt_rhs=e'}
= (newAlt, top_fvs, lcl_fvs')