%
% (c) The GRASP/AQUA Project, Glasgow University, 19921998
%
\section{SetLevels}
***************************
Overview
***************************
1. We attach binding levels to Core bindings, in preparation for floating
outwards (@FloatOut@).
2. We also letify many expressions (notably case scrutinees), so they
will have a fighting chance of being floated sensible.
3. We clone the binders of any floatable letbinding, so that when it is
floated out it will be unique. (This used to be done by the simplifier
but the latter now only ensures that there's no shadowing; indeed, even
that may not be true.)
NOTE: this can't be done using the uniqAway idea, because the variable
must be unique in the whole program, not just its current scope,
because two variables in different scopes may float out to the
same top level place
NOTE: Very tiresomely, we must apply this substitution to
the rules stored inside a variable too.
We do *not* clone toplevel bindings, because some of them must not change,
but we *do* clone bindings that are heading for the top level
4. In the expression
case x of wild { p -> ...wild... }
we substitute x for wild in the RHS of the case alternatives:
case x of wild { p -> ...x... }
This means that a subexpression involving x is not "trapped" inside the RHS.
And it's not inconvenient because we already have a substitution.
Note that this is EXACTLY BACKWARDS from the what the simplifier does.
The simplifier tries to get rid of occurrences of x, in favour of wild,
in the hope that there will only be one remaining occurrence of x, namely
the scrutinee of the case, and we can inline it.
\begin{code}
module SetLevels (
setLevels,
Level(..), tOP_LEVEL,
LevelledBind, LevelledExpr,
incMinorLvl, ltMajLvl, ltLvl, isTopLvl, isInlineCtxt
) where
#include "HsVersions.h"
import CoreSyn
import DynFlags ( FloatOutSwitches(..) )
import CoreUtils ( exprType, exprIsTrivial, mkPiTypes )
import CoreFVs
import CoreSubst ( Subst, emptySubst, extendInScope, extendIdSubst,
cloneIdBndr, cloneRecIdBndrs )
import Id ( idType, mkSysLocal, isOneShotLambda,
zapDemandIdInfo, transferPolyIdInfo,
idSpecialisation, idWorkerInfo, setIdInfo
)
import IdInfo
import Var
import VarSet
import VarEnv
import Name ( getOccName )
import OccName ( occNameString )
import Type ( isUnLiftedType, Type )
import BasicTypes ( TopLevelFlag(..) )
import UniqSupply
import Util ( sortLe, isSingleton, count )
import Outputable
import FastString
\end{code}
%************************************************************************
%* *
\subsection{Level numbers}
%* *
%************************************************************************
\begin{code}
data Level = InlineCtxt
| Level Int
Int
\end{code}
The {\em level number} on a (type)lambdabound variable is the
nesting depth of the (type)lambda which binds it. The outermost lambda
has level 1, so (Level 0 0) means that the variable is bound outside any lambda.
On an expression, it's the maximum level number of its free
(type)variables. On a let(rec)bound variable, it's the level of its
RHS. On a casebound variable, it's the number of enclosing lambdas.
Toplevel variables: level~0. Those bound on the RHS of a toplevel
definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
as ``subscripts'')...
\begin{verbatim}
a_0 = let b_? = ... in
x_1 = ... b ... in ...
\end{verbatim}
The main function @lvlExpr@ carries a ``context level'' (@ctxt_lvl@).
That's meant to be the level number of the enclosing binder in the
final (floated) program. If the level number of a subexpression is
less than that of the context, then it might be worth letbinding the
subexpression so that it will indeed float.
If you can float to level @Level 0 0@ worth doing so because then your
allocation becomes static instead of dynamic. We always start with
context @Level 0 0@.
Note [FloatOut inside INLINE]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@InlineCtxt@ very similar to @Level 0 0@, but is used for one purpose:
to say "don't float anything out of here". That's exactly what we
want for the body of an INLINE, where we don't want to float anything
out at all. See notes with lvlMFE below.
But, check this out:
So the conclusion is: don't do any floating at all inside an InlineMe.
(In the above example, don't float the {x=e} out of the \d.)
One particular case is that of workers: we don't want to float the
call to the worker outside the wrapper, otherwise the worker might get
inlined into the floated expression, and an importing module won't see
the worker at all.
\begin{code}
type LevelledExpr = TaggedExpr Level
type LevelledBind = TaggedBind Level
tOP_LEVEL, iNLINE_CTXT :: Level
tOP_LEVEL = Level 0 0
iNLINE_CTXT = InlineCtxt
incMajorLvl :: Level -> Level
incMajorLvl InlineCtxt = InlineCtxt
incMajorLvl (Level major _) = Level (major + 1) 0
incMinorLvl :: Level -> Level
incMinorLvl InlineCtxt = InlineCtxt
incMinorLvl (Level major minor) = Level major (minor+1)
maxLvl :: Level -> Level -> Level
maxLvl InlineCtxt l2 = l2
maxLvl l1 InlineCtxt = l1
maxLvl l1@(Level maj1 min1) l2@(Level maj2 min2)
| (maj1 > maj2) || (maj1 == maj2 && min1 > min2) = l1
| otherwise = l2
ltLvl :: Level -> Level -> Bool
ltLvl _ InlineCtxt = False
ltLvl InlineCtxt (Level _ _) = True
ltLvl (Level maj1 min1) (Level maj2 min2)
= (maj1 < maj2) || (maj1 == maj2 && min1 < min2)
ltMajLvl :: Level -> Level -> Bool
ltMajLvl _ InlineCtxt = False
ltMajLvl InlineCtxt (Level maj2 _) = 0 < maj2
ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2
isTopLvl :: Level -> Bool
isTopLvl (Level 0 0) = True
isTopLvl _ = False
isInlineCtxt :: Level -> Bool
isInlineCtxt InlineCtxt = True
isInlineCtxt _ = False
instance Outputable Level where
ppr InlineCtxt = text "<INLINE>"
ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]
instance Eq Level where
InlineCtxt == InlineCtxt = True
(Level maj1 min1) == (Level maj2 min2) = maj1 == maj2 && min1 == min2
_ == _ = False
\end{code}
%************************************************************************
%* *
\subsection{Main levelsetting code}
%* *
%************************************************************************
\begin{code}
setLevels :: FloatOutSwitches
-> [CoreBind]
-> UniqSupply
-> [LevelledBind]
setLevels float_lams binds us
= initLvl us (do_them binds)
where
do_them :: [CoreBind] -> LvlM [LevelledBind]
do_them [] = return []
do_them (b:bs) = do
(lvld_bind, _) <- lvlTopBind init_env b
lvld_binds <- do_them bs
return (lvld_bind : lvld_binds)
init_env = initialEnv float_lams
lvlTopBind :: LevelEnv -> Bind Id -> LvlM (LevelledBind, LevelEnv)
lvlTopBind env (NonRec binder rhs)
= lvlBind TopLevel tOP_LEVEL env (AnnNonRec binder (freeVars rhs))
lvlTopBind env (Rec pairs)
= lvlBind TopLevel tOP_LEVEL env (AnnRec [(b,freeVars rhs) | (b,rhs) <- pairs])
\end{code}
%************************************************************************
%* *
\subsection{Setting expression levels}
%* *
%************************************************************************
\begin{code}
lvlExpr :: Level
-> LevelEnv
-> CoreExprWithFVs
-> LvlM LevelledExpr
\end{code}
The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
binder. Here's an example
v = \x -> ...\y -> let r = case (..x..) of
..x..
in ..
When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
the level of @r@, even though it's inside a level2 @\y@. It's
important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
If there were another lambda in @r@'s rhs, it would get level2 as well.
\begin{code}
lvlExpr _ _ ( _, AnnType ty) = return (Type ty)
lvlExpr _ env (_, AnnVar v) = return (lookupVar env v)
lvlExpr _ _ (_, AnnLit lit) = return (Lit lit)
lvlExpr ctxt_lvl env (_, AnnApp fun arg) = do
fun' <- lvl_fun fun
arg' <- lvlMFE False ctxt_lvl env arg
return (App fun' arg')
where
lvl_fun (_, AnnCase _ _ _ _) = lvlMFE True ctxt_lvl env fun
lvl_fun _ = lvlExpr ctxt_lvl env fun
lvlExpr _ env (_, AnnNote InlineMe expr) = do
expr' <- lvlExpr iNLINE_CTXT env expr
return (Note InlineMe expr')
lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
expr' <- lvlExpr ctxt_lvl env expr
return (Note note expr')
lvlExpr ctxt_lvl env (_, AnnCast expr co) = do
expr' <- lvlExpr ctxt_lvl env expr
return (Cast expr' co)
lvlExpr ctxt_lvl env expr@(_, AnnLam {}) = do
new_body <- lvlMFE True new_lvl new_env body
return (mkLams new_bndrs new_body)
where
(bndrs, body) = collectAnnBndrs expr
(new_lvl, new_bndrs) = lvlLamBndrs ctxt_lvl bndrs
new_env = extendLvlEnv env new_bndrs
lvlExpr ctxt_lvl env (_, AnnLet (AnnNonRec bndr rhs) body)
| isUnLiftedType (idType bndr) = do
rhs' <- lvlExpr ctxt_lvl env rhs
body' <- lvlExpr incd_lvl env' body
return (Let (NonRec bndr' rhs') body')
where
incd_lvl = incMinorLvl ctxt_lvl
bndr' = TB bndr incd_lvl
env' = extendLvlEnv env [bndr']
lvlExpr ctxt_lvl env (_, AnnLet bind body) = do
(bind', new_env) <- lvlBind NotTopLevel ctxt_lvl env bind
body' <- lvlExpr ctxt_lvl new_env body
return (Let bind' body')
lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr ty alts) = do
expr' <- lvlMFE True ctxt_lvl env expr
let alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
alts' <- mapM (lvl_alt alts_env) alts
return (Case expr' (TB case_bndr incd_lvl) ty alts')
where
incd_lvl = incMinorLvl ctxt_lvl
lvl_alt alts_env (con, bs, rhs) = do
rhs' <- lvlMFE True incd_lvl new_env rhs
return (con, bs', rhs')
where
bs' = [ TB b incd_lvl | b <- bs ]
new_env = extendLvlEnv alts_env bs'
\end{code}
@lvlMFE@ is just like @lvlExpr@, except that it might letbind
the expression, so that it can itself be floated.
Note [Unlifted MFEs]
~~~~~~~~~~~~~~~~~~~~~
We don't float unlifted MFEs, which potentially loses big opportunites.
For example:
\x -> f (h y)
where h :: Int -> Int# is expensive. We'd like to float the (h y) outside
the \x, but we don't because it's unboxed. Possible solution: box it.
Note [Case MFEs]
~~~~~~~~~~~~~~~~
We don't float a case expression as an MFE from a strict context. Why not?
Because in doing so we share a tiny bit of computation (the switch) but
in exchange we build a thunk, which is bad. This case reduces allocation
by 7% in spectral/puzzle (a rather strange benchmark) and 1.2% in real/fem.
Doesn't change any other allocation at all.
\begin{code}
lvlMFE :: Bool
-> Level
-> LevelEnv
-> CoreExprWithFVs
-> LvlM LevelledExpr
lvlMFE _ _ _ (_, AnnType ty)
= return (Type ty)
lvlMFE strict_ctxt ctxt_lvl env (_, AnnCast e co)
= do { expr' <- lvlMFE strict_ctxt ctxt_lvl env e
; return (Cast expr' co) }
lvlMFE True ctxt_lvl env e@(_, AnnCase {})
= lvlExpr ctxt_lvl env e
lvlMFE strict_ctxt ctxt_lvl env ann_expr@(fvs, _)
| isUnLiftedType ty
|| isInlineCtxt ctxt_lvl
|| exprIsTrivial expr
|| not good_destination
=
lvlExpr ctxt_lvl env ann_expr
| otherwise
= do expr' <- lvlFloatRhs abs_vars dest_lvl env ann_expr
var <- newLvlVar "lvl" abs_vars ty
return (Let (NonRec (TB var dest_lvl) expr')
(mkVarApps (Var var) abs_vars))
where
expr = deAnnotate ann_expr
ty = exprType expr
dest_lvl = destLevel env fvs (isFunction ann_expr)
abs_vars = abstractVars dest_lvl env fvs
good_destination
| dest_lvl `ltMajLvl` ctxt_lvl
= True
| otherwise
= isTopLvl dest_lvl
&& floatConsts env
&& not strict_ctxt
\end{code}
Note [Escaping a value lambda]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We want to float even cheap expressions out of value lambdas,
because that saves allocation. Consider
f = \x. .. (\y.e) ...
Then we'd like to avoid allocating the (\y.e) every time we call f,
(assuming e does not mention x).
An example where this really makes a difference is simplrun009.
Another reason it's good is because it makes SpecContr fire on functions.
Consider
f = \x. ....(f (\y.e))....
After floating we get
lvl = \y.e
f = \x. ....(f lvl)...
and that is much easier for SpecConstr to generate a robust specialisation for.
The OLD CODE (given where this Note is referred to) prevents floating
of the example above, so I just don't understand the old code. I
don't understand the old comment either (which appears below). I
measured the effect on nofib of changing OLD CODE to 'True', and got
zeros everywhere, but a 4% win for 'puzzle'. Very small 0.5% loss for
'cse'; turns out to be because our arity analysis isn't good enough
yet (mentioned in Simonnofibnotes).
OLD comment was:
Even if it escapes a value lambda, we only
float if it's not cheap (unless it'll get all the
way to the top). I've seen cases where we
float dozens of tiny free expressions, which cost
more to allocate than to evaluate.
NB: exprIsCheap is also true of bottom expressions, which
is good; we don't want to share them
It's only Really Bad to float a cheap expression out of a
strict context, because that builds a thunk that otherwise
would never be built. So another alternative would be to
add
|| (strict_ctxt && not (exprIsBottom expr))
to the condition above. We should really try this out.
%************************************************************************
%* *
\subsection{Bindings}
%* *
%************************************************************************
The binding stuff works for top level too.
\begin{code}
lvlBind :: TopLevelFlag
-> Level
-> LevelEnv
-> CoreBindWithFVs
-> LvlM (LevelledBind, LevelEnv)
lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
| isTyVar bndr
|| isInlineCtxt ctxt_lvl
= do rhs' <- lvlExpr ctxt_lvl env rhs
return (NonRec (TB bndr ctxt_lvl) rhs', env)
| null abs_vars
= do
rhs' <- lvlExpr dest_lvl env rhs
(env', bndr') <- cloneVar top_lvl env bndr ctxt_lvl dest_lvl
return (NonRec (TB bndr' dest_lvl) rhs', env')
| otherwise
= do
rhs' <- lvlFloatRhs abs_vars dest_lvl env rhs
(env', [bndr']) <- newPolyBndrs dest_lvl env abs_vars [bndr]
return (NonRec (TB bndr' dest_lvl) rhs', env')
where
bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
abs_vars = abstractVars dest_lvl env bind_fvs
dest_lvl = destLevel env bind_fvs (isFunction rhs)
\end{code}
\begin{code}
lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
| isInlineCtxt ctxt_lvl
= do rhss' <- mapM (lvlExpr ctxt_lvl env) rhss
return (Rec ([TB b ctxt_lvl | b <- bndrs] `zip` rhss'), env)
| null abs_vars
= do (new_env, new_bndrs) <- cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl
new_rhss <- mapM (lvlExpr ctxt_lvl new_env) rhss
return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
| isSingleton pairs && count isId abs_vars > 1
= do
let
(bndr,rhs) = head pairs
(rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
rhs_env = extendLvlEnv env abs_vars_w_lvls
(rhs_env', new_bndr) <- cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl
let
(lam_bndrs, rhs_body) = collectAnnBndrs rhs
(body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
body_env = extendLvlEnv rhs_env' new_lam_bndrs
new_rhs_body <- lvlExpr body_lvl body_env rhs_body
(poly_env, [poly_bndr]) <- newPolyBndrs dest_lvl env abs_vars [bndr]
return (Rec [(TB poly_bndr dest_lvl,
mkLams abs_vars_w_lvls $
mkLams new_lam_bndrs $
Let (Rec [(TB new_bndr rhs_lvl, mkLams new_lam_bndrs new_rhs_body)])
(mkVarApps (Var new_bndr) lam_bndrs))],
poly_env)
| otherwise = do
(new_env, new_bndrs) <- newPolyBndrs dest_lvl env abs_vars bndrs
new_rhss <- mapM (lvlFloatRhs abs_vars dest_lvl new_env) rhss
return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
where
(bndrs,rhss) = unzip pairs
bind_fvs = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
| (bndr, (rhs_fvs,_)) <- pairs])
`minusVarSet`
mkVarSet bndrs
dest_lvl = destLevel env bind_fvs (all isFunction rhss)
abs_vars = abstractVars dest_lvl env bind_fvs
lvlFloatRhs :: [CoreBndr] -> Level -> LevelEnv -> CoreExprWithFVs
-> UniqSM (Expr (TaggedBndr Level))
lvlFloatRhs abs_vars dest_lvl env rhs = do
rhs' <- lvlExpr rhs_lvl rhs_env rhs
return (mkLams abs_vars_w_lvls rhs')
where
(rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
rhs_env = extendLvlEnv env abs_vars_w_lvls
\end{code}
%************************************************************************
%* *
\subsection{Deciding floatability}
%* *
%************************************************************************
\begin{code}
lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [TaggedBndr Level])
lvlLamBndrs lvl []
= (lvl, [])
lvlLamBndrs lvl bndrs
= go (incMinorLvl lvl)
False
[] bndrs
where
go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
| isId bndr &&
not bumped_major &&
not (isOneShotLambda bndr)
= go new_lvl True (TB bndr new_lvl : rev_lvld_bndrs) bndrs
| otherwise
= go old_lvl bumped_major (TB bndr old_lvl : rev_lvld_bndrs) bndrs
where
new_lvl = incMajorLvl old_lvl
go old_lvl _ rev_lvld_bndrs []
= (old_lvl, reverse rev_lvld_bndrs)
\end{code}
\begin{code}
destLevel :: LevelEnv -> VarSet -> Bool -> Level
destLevel env fvs is_function
| floatLams env
&& is_function = tOP_LEVEL
| otherwise = maxIdLevel env fvs
isFunction :: CoreExprWithFVs -> Bool
isFunction (_, AnnLam b e) | isId b = True
| otherwise = isFunction e
isFunction (_, AnnNote _ e) = isFunction e
isFunction _ = False
\end{code}
%************************************************************************
%* *
\subsection{FreeToLevel Monad}
%* *
%************************************************************************
\begin{code}
type LevelEnv = (FloatOutSwitches,
VarEnv Level,
Subst,
IdEnv ([Var], LevelledExpr))
initialEnv :: FloatOutSwitches -> LevelEnv
initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
floatLams :: LevelEnv -> Bool
floatLams (fos, _, _, _) = floatOutLambdas fos
floatConsts :: LevelEnv -> Bool
floatConsts (fos, _, _, _) = floatOutConstants fos
extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
= (float_lams,
foldl add_lvl lvl_env prs,
foldl del_subst subst prs,
foldl del_id id_env prs)
where
add_lvl env (TB v l) = extendVarEnv env v l
del_subst env (TB v _) = extendInScope env v
del_id env (TB v _) = delVarEnv env v
extendCaseBndrLvlEnv :: LevelEnv -> Expr (TaggedBndr Level) -> Var -> Level
-> LevelEnv
extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
= (float_lams,
extendVarEnv lvl_env case_bndr lvl,
extendIdSubst subst case_bndr (Var scrut_var),
extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
extendCaseBndrLvlEnv env _scrut case_bndr lvl
= extendLvlEnv env [TB case_bndr lvl]
extendPolyLvlEnv :: Level -> LevelEnv -> [Var] -> [(Var, Var)] -> LevelEnv
extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
= (float_lams,
foldl add_lvl lvl_env bndr_pairs,
foldl add_subst subst bndr_pairs,
foldl add_id id_env bndr_pairs)
where
add_lvl env (_, v') = extendVarEnv env v' dest_lvl
add_subst env (v, v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
add_id env (v, v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
extendCloneLvlEnv :: Level -> LevelEnv -> Subst -> [(Var, Var)] -> LevelEnv
extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
= (float_lams,
foldl add_lvl lvl_env bndr_pairs,
new_subst,
foldl add_id id_env bndr_pairs)
where
add_lvl env (_, v') = extendVarEnv env v' lvl
add_id env (v, v') = extendVarEnv env v ([v'], Var v')
maxIdLevel :: LevelEnv -> VarSet -> Level
maxIdLevel (_, lvl_env,_,id_env) var_set
= foldVarSet max_in tOP_LEVEL var_set
where
max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
Just (abs_vars, _) -> abs_vars
Nothing -> [in_var])
max_out out_var lvl
| isId out_var = case lookupVarEnv lvl_env out_var of
Just lvl' -> maxLvl lvl' lvl
Nothing -> lvl
| otherwise = lvl
lookupVar :: LevelEnv -> Id -> LevelledExpr
lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
Just (_, expr) -> expr
_ -> Var v
abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
abstractVars dest_lvl (_, lvl_env, _, id_env) fvs
= map zap $ uniq $ sortLe le
[var | fv <- varSetElems fvs
, var <- absVarsOf id_env fv
, abstract_me var ]
where
v1 `le` v2 = case (is_tv v1, is_tv v2) of
(True, False) -> True
(False, True) -> False
_ -> v1 <= v2
is_tv v = isTyVar v && not (isCoVar v)
uniq :: [Var] -> [Var]
uniq (v1:v2:vs) | v1 == v2 = uniq (v2:vs)
| otherwise = v1 : uniq (v2:vs)
uniq vs = vs
abstract_me v = case lookupVarEnv lvl_env v of
Just lvl -> dest_lvl `ltLvl` lvl
Nothing -> False
zap v | isId v = WARN( workerExists (idWorkerInfo v) ||
not (isEmptySpecInfo (idSpecialisation v)),
text "absVarsOf: discarding info on" <+> ppr v )
setIdInfo v vanillaIdInfo
| otherwise = v
absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> [Var]
absVarsOf id_env v
| isId v = [av2 | av1 <- lookup_avs v
, av2 <- add_tyvars av1]
| isCoVar v = add_tyvars v
| otherwise = [v]
where
lookup_avs v = case lookupVarEnv id_env v of
Just (abs_vars, _) -> abs_vars
Nothing -> [v]
add_tyvars v = v : varSetElems (varTypeTyVars v)
\end{code}
\begin{code}
type LvlM result = UniqSM result
initLvl :: UniqSupply -> UniqSM a -> a
initLvl = initUs_
\end{code}
\begin{code}
newPolyBndrs :: Level -> LevelEnv -> [Var] -> [Id] -> UniqSM (LevelEnv, [Id])
newPolyBndrs dest_lvl env abs_vars bndrs = do
uniqs <- getUniquesM
let new_bndrs = zipWith mk_poly_bndr bndrs uniqs
return (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
where
mk_poly_bndr bndr uniq = transferPolyIdInfo bndr abs_vars $
mkSysLocal (mkFastString str) uniq poly_ty
where
str = "poly_" ++ occNameString (getOccName bndr)
poly_ty = mkPiTypes abs_vars (idType bndr)
newLvlVar :: String
-> [CoreBndr] -> Type
-> LvlM Id
newLvlVar str vars body_ty = do
uniq <- getUniqueM
return (mkSysLocal (mkFastString str) uniq (mkPiTypes vars body_ty))
cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
cloneVar TopLevel env v _ _
= return (env, v)
cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
= ASSERT( isId v ) do
us <- getUniqueSupplyM
let
(subst', v1) = cloneIdBndr subst us v
v2 = zap_demand ctxt_lvl dest_lvl v1
env' = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
return (env', v2)
cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
cloneRecVars TopLevel env vs _ _
= return (env, vs)
cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
= ASSERT( all isId vs ) do
us <- getUniqueSupplyM
let
(subst', vs1) = cloneRecIdBndrs subst us vs
vs2 = map (zap_demand ctxt_lvl dest_lvl) vs1
env' = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
return (env', vs2)
zap_demand :: Level -> Level -> Id -> Id
zap_demand dest_lvl ctxt_lvl id
| ctxt_lvl == dest_lvl,
not (isTopLvl dest_lvl) = id
| otherwise = zapDemandIdInfo id
\end{code}