%
-% (c) The AQUA Project, Glasgow University, 1994-1996
+% (c) The AQUA Project, Glasgow University, 1994-1998
%
\section[LambdaLift]{A STG-code lambda lifter}
\begin{code}
-#include "HsVersions.h"
-
module LambdaLift ( liftProgram ) where
-import Ubiq{-uitous-}
+#include "HsVersions.h"
import StgSyn
-import Bag ( emptyBag, unionBags, unitBag, snocBag, bagToList )
-import Id ( idType, mkSysLocal, addIdArity,
- mkIdSet, unitIdSet, minusIdSet,
- unionManyIdSets, idSetToList, IdSet(..),
- nullIdEnv, growIdEnvList, lookupIdEnv, IdEnv(..)
- )
-import SrcLoc ( mkUnknownSrcLoc )
-import Type ( splitForAllTy, mkForAllTys, mkFunTys )
-import UniqSupply ( getUnique, splitUniqSupply )
-import Util ( zipEqual, panic, assertPanic )
+import Bag ( Bag, emptyBag, unionBags, unitBag, snocBag, bagToList )
+import Id ( mkVanillaId, idType, setIdArityInfo, Id )
+import VarSet
+import VarEnv
+import IdInfo ( exactArity )
+import Module ( Module )
+import Name ( mkTopName )
+import Type ( splitForAllTys, mkForAllTys, mkFunTys, Type )
+import UniqSupply ( uniqFromSupply, splitUniqSupply, UniqSupply )
+import Util ( zipEqual )
+import Panic ( panic, assertPanic )
\end{code}
This is the lambda lifter. It turns lambda abstractions into
recursive calls, which may now have lots of free vars.
Recent Observations:
+
* 2 might be already ``too many'' variables to abstract.
The problem is that the increase in the number of free variables
of closures refering to the lifted function (which is always # of
abstracted args - 1) may increase heap allocation a lot.
Expeiments are being done to check this...
+
* We do not lambda lift if the function has at least one occurrence
without any arguments. This caused lots of problems. Ex:
h = \ x -> ... let y = ...
%************************************************************************
\begin{code}
-liftProgram :: UniqSupply -> [StgBinding] -> [StgBinding]
-liftProgram us prog = concat (runLM Nothing us (mapLM liftTopBind prog))
+liftProgram :: Module -> UniqSupply -> [StgBinding] -> [StgBinding]
+liftProgram mod us prog = concat (runLM mod Nothing us (mapLM liftTopBind prog))
liftTopBind :: StgBinding -> LiftM [StgBinding]
-> LiftM (StgExpr, LiftInfo)
-liftExpr expr@(StgCon con args lvs) = returnLM (expr, emptyLiftInfo)
-liftExpr expr@(StgPrim op args lvs) = returnLM (expr, emptyLiftInfo)
+liftExpr expr@(StgLit _) = returnLM (expr, emptyLiftInfo)
+liftExpr expr@(StgConApp _ _) = returnLM (expr, emptyLiftInfo)
+liftExpr expr@(StgPrimApp _ _ _) = returnLM (expr, emptyLiftInfo)
-liftExpr expr@(StgApp (StgLitArg lit) args lvs) = returnLM (expr, emptyLiftInfo)
-liftExpr expr@(StgApp (StgVarArg v) args lvs)
- = lookup v `thenLM` \ ~(sc, sc_args) -> -- NB the ~. We don't want to
+liftExpr expr@(StgApp v args)
+ = lookUp v `thenLM` \ ~(sc, sc_args) -> -- NB the ~. We don't want to
-- poke these bindings too early!
- returnLM (StgApp (StgVarArg sc) (map StgVarArg sc_args ++ args) lvs,
+ returnLM (StgApp sc (map StgVarArg sc_args ++ args),
emptyLiftInfo)
-- The lvs field is probably wrong, but we reconstruct it
-- anyway following lambda lifting
-liftExpr (StgCase scrut lv1 lv2 uniq alts)
+liftExpr (StgCase scrut lv1 lv2 bndr srt alts)
= liftExpr scrut `thenLM` \ (scrut', scrut_info) ->
lift_alts alts `thenLM` \ (alts', alts_info) ->
- returnLM (StgCase scrut' lv1 lv2 uniq alts', scrut_info `unionLiftInfo` alts_info)
+ returnLM (StgCase scrut' lv1 lv2 bndr srt alts', scrut_info `unionLiftInfo` alts_info)
where
lift_alts (StgAlgAlts ty alg_alts deflt)
= mapAndUnzipLM lift_alg_alt alg_alts `thenLM` \ (alg_alts', alt_infos) ->
returnLM ((lit, rhs'), rhs_info)
lift_deflt StgNoDefault = returnLM (StgNoDefault, emptyLiftInfo)
- lift_deflt (StgBindDefault var used rhs)
+ lift_deflt (StgBindDefault rhs)
= liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
- returnLM (StgBindDefault var used rhs', rhs_info)
+ returnLM (StgBindDefault rhs', rhs_info)
\end{code}
Now the interesting cases. Let no escape isn't lifted. We turn it
let
-- Find the free vars of all the rhss,
-- excluding the binders themselves.
- rhs_free_vars = unionManyIdSets (map rhsFreeVars rhss)
- `minusIdSet`
- mkIdSet binders
+ rhs_free_vars = unionVarSets (map rhsFreeVars rhss)
+ `minusVarSet`
+ mkVarSet binders
rhs_info = unionLiftInfos rhs_infos
in
\end{code}
\begin{code}
-liftExpr (StgSCC ty cc expr)
+liftExpr (StgSCC cc expr)
= liftExpr expr `thenLM` \ (expr2, expr_info) ->
- returnLM (StgSCC ty cc expr2, expr_info)
+ returnLM (StgSCC cc expr2, expr_info)
\end{code}
A binding is liftable if it's a *function* (args not null) and never
\begin{code}
isLiftable :: StgRhs -> Bool
-isLiftable (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
+isLiftable (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) _ fvs _ args _)
-- Experimental evidence suggests we should lift only if we will be
-- abstracting up to 4 fvs.
here).
-}
-isLiftableRec (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
+isLiftableRec (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) _ fvs _ args _)
= if not (null args || -- Not a function
unapplied_occ || -- Has an occ with no args at all
arg_occ || -- Occurs in arg position
isLiftableRec other_rhs = False
rhsFreeVars :: StgRhs -> IdSet
-rhsFreeVars (StgRhsClosure _ _ fvs _ _ _) = mkIdSet fvs
+rhsFreeVars (StgRhsClosure _ _ _ fvs _ _ _) = mkVarSet fvs
rhsFreeVars other = panic "rhsFreeVars"
\end{code}
dontLiftRhs rhs@(StgRhsCon cc v args) = returnLM (rhs, emptyLiftInfo)
-dontLiftRhs (StgRhsClosure cc bi fvs upd args body)
+dontLiftRhs (StgRhsClosure cc bi srt fvs upd args body)
= liftExpr body `thenLM` \ (body', body_info) ->
- returnLM (StgRhsClosure cc bi fvs upd args body', body_info)
+ returnLM (StgRhsClosure cc bi srt fvs upd args body', body_info)
\end{code}
\begin{code}
-- the set is its free vars
(Id,StgRhs)) -- Binding for supercombinator
-mkScPieces extra_arg_set (id, StgRhsClosure cc bi _ upd args body)
+mkScPieces extra_arg_set (id, StgRhsClosure cc bi srt _ upd args body)
= ASSERT( n_args > 0 )
-- Construct the rhs of the supercombinator, and its Id
newSupercombinator sc_ty arity `thenLM` \ sc_id ->
returnLM ((sc_id, extra_args), (sc_id, sc_rhs))
where
n_args = length args
- extra_args = idSetToList extra_arg_set
+ extra_args = varSetElems extra_arg_set
arity = n_args + length extra_args
-- Construct the supercombinator type
type_of_original_id = idType id
extra_arg_tys = map idType extra_args
- (tyvars, rest) = splitForAllTy type_of_original_id
+ (tyvars, rest) = splitForAllTys type_of_original_id
sc_ty = mkForAllTys tyvars (mkFunTys extra_arg_tys rest)
- sc_rhs = StgRhsClosure cc bi [] upd (extra_args ++ args) body
+ sc_rhs = StgRhsClosure cc bi srt [] upd (extra_args ++ args) body
\end{code}
The monad is used only to distribute global stuff, and the unique supply.
\begin{code}
-type LiftM a = LiftFlags
+type LiftM a = Module
+ -> LiftFlags
-> UniqSupply
-> (IdEnv -- Domain = candidates for lifting
(Id, -- The supercombinator
-- binding; Nothing == infinity
-runLM :: LiftFlags -> UniqSupply -> LiftM a -> a
-runLM flags us m = m flags us nullIdEnv
+runLM :: Module -> LiftFlags -> UniqSupply -> LiftM a -> a
+runLM mod flags us m = m mod flags us emptyVarEnv
thenLM :: LiftM a -> (a -> LiftM b) -> LiftM b
-thenLM m k ci us idenv
- = k (m ci us1 idenv) ci us2 idenv
+thenLM m k mod ci us idenv
+ = k (m mod ci us1 idenv) mod ci us2 idenv
where
(us1, us2) = splitUniqSupply us
returnLM :: a -> LiftM a
-returnLM a ci us idenv = a
+returnLM a mod ci us idenv = a
fixLM :: (a -> LiftM a) -> LiftM a
-fixLM k ci us idenv = r
+fixLM k mod ci us idenv = r
where
- r = k r ci us idenv
+ r = k r mod ci us idenv
mapLM :: (a -> LiftM b) -> [a] -> LiftM [b]
mapLM f [] = returnLM []
-> Int -- Arity
-> LiftM Id
-newSupercombinator ty arity ci us idenv
- = (mkSysLocal SLIT("sc") uniq ty mkUnknownSrcLoc) -- ToDo: improve location
- `addIdArity` arity
- -- ToDo: rm the addIdArity? Just let subsequent stg-saturation pass do it?
+newSupercombinator ty arity mod ci us idenv
+ = mkVanillaId (mkTopName uniq mod SLIT("_ll")) ty
+ `setIdArityInfo` exactArity arity
+ -- ToDo: rm the setIdArity? Just let subsequent stg-saturation pass do it?
where
- uniq = getUnique us
+ uniq = uniqFromSupply us
-lookup :: Id -> LiftM (Id,[Id])
-lookup v ci us idenv
- = case (lookupIdEnv idenv v) of
+lookUp :: Id -> LiftM (Id,[Id])
+lookUp v mod ci us idenv
+ = case (lookupVarEnv idenv v) of
Just result -> result
Nothing -> (v, [])
addScInlines :: [Id] -> [(Id,[Id])] -> LiftM a -> LiftM a
-addScInlines ids values m ci us idenv
- = m ci us idenv'
+addScInlines ids values m mod ci us idenv
+ = m mod ci us idenv'
where
- idenv' = growIdEnvList idenv (ids `zip_lazy` values)
+ idenv' = extendVarEnvList idenv (ids `zip_lazy` values)
-- zip_lazy zips two things together but matches lazily on the
-- second argument. This is important, because the ids are know here,
getFinalFreeVars :: IdSet -> LiftM IdSet
-getFinalFreeVars free_vars ci us idenv
- = unionManyIdSets (map munge_it (idSetToList free_vars))
+getFinalFreeVars free_vars mod ci us idenv
+ = unionVarSets (map munge_it (varSetElems free_vars))
where
munge_it :: Id -> IdSet -- Takes a free var and maps it to the "real"
-- free var
- munge_it id = case (lookupIdEnv idenv id) of
- Just (_, args) -> mkIdSet args
- Nothing -> unitIdSet id
+ munge_it id = case (lookupVarEnv idenv id) of
+ Just (_, args) -> mkVarSet args
+ Nothing -> unitVarSet id
\end{code}
getScBinds :: LiftInfo -> [StgBinding]
getScBinds binds = bagToList binds
-looksLikeSATRhs [(f,StgRhsClosure _ _ _ _ ls _)] (StgApp (StgVarArg f') args _)
+looksLikeSATRhs [(f,StgRhsClosure _ _ _ _ _ ls _)] (StgApp f' args)
= (f == f') && (length args == length ls)
looksLikeSATRhs _ _ = False
\end{code}