2 % (c) The AQUA Project, Glasgow University, 1994-1996
4 \section[LambdaLift]{A STG-code lambda lifter}
7 #include "HsVersions.h"
9 module LambdaLift ( liftProgram ) where
15 import Bag ( emptyBag, unionBags, unitBag, snocBag, bagToList )
16 import Id ( idType, mkSysLocal, addIdArity,
17 mkIdSet, unitIdSet, minusIdSet,
18 unionManyIdSets, idSetToList, IdSet(..),
19 nullIdEnv, growIdEnvList, lookupIdEnv, IdEnv(..)
21 import SrcLoc ( mkUnknownSrcLoc )
22 import Type ( splitForAllTy, mkForAllTys, mkFunTys )
23 import UniqSupply ( getUnique, splitUniqSupply )
24 import Util ( zipEqual, panic, assertPanic )
27 This is the lambda lifter. It turns lambda abstractions into
28 supercombinators on a selective basis:
30 * Let-no-escaped bindings are never lifted. That's one major reason
31 why the lambda lifter is done in STG.
33 * Non-recursive bindings whose RHS is a lambda abstractions are lifted,
34 provided all the occurrences of the bound variable is in a function
35 postition. In this example, f will be lifted:
43 $f p q r x = e -- Supercombinator
45 ..($f p q r a1)...($f p q r a2)...
47 NOTE that the original binding is eliminated.
49 But in this case, f won't be lifted:
56 Why? Because we have to heap-allocate a closure for f thus:
58 $f p q r x = e -- Supercombinator
63 ..(g f)...($f p q r a2)..
65 so it might as well be the original lambda abstraction.
67 We also do not lift if the function has an occurrence with no arguments, e.g.
73 as this form is more efficient than if we create a partial application
75 $f p q r x = e -- Supercombinator
79 * Recursive bindings *all* of whose RHSs are lambda abstractions are
81 - all the occurrences of all the binders are in a function position
82 - there aren't ``too many'' free variables.
84 Same reasoning as before for the function-position stuff. The ``too many
85 free variable'' part comes from considering the (potentially many)
86 recursive calls, which may now have lots of free vars.
89 * 2 might be already ``too many'' variables to abstract.
90 The problem is that the increase in the number of free variables
91 of closures refering to the lifted function (which is always # of
92 abstracted args - 1) may increase heap allocation a lot.
93 Expeiments are being done to check this...
94 * We do not lambda lift if the function has at least one occurrence
95 without any arguments. This caused lots of problems. Ex:
96 h = \ x -> ... let y = ...
97 in let let f = \x -> ...y...
101 h = \ x -> ... let y = ...
104 now f y is a partial application, so it will be updated, and this
108 --- NOT RELEVANT FOR STG ----
109 * All ``lone'' lambda abstractions are lifted. Notably this means lambda
111 - in a case alternative: case e of True -> (\x->b)
112 - in the body of a let: let x=e in (\y->b)
113 -----------------------------
115 %************************************************************************
117 \subsection[Lift-expressions]{The main function: liftExpr}
119 %************************************************************************
122 liftProgram :: UniqSupply -> [StgBinding] -> [StgBinding]
123 liftProgram us prog = concat (runLM Nothing us (mapLM liftTopBind prog))
126 liftTopBind :: StgBinding -> LiftM [StgBinding]
127 liftTopBind (StgNonRec id rhs)
128 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
129 returnLM (getScBinds rhs_info ++ [StgNonRec id rhs'])
131 liftTopBind (StgRec pairs)
132 = mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
133 returnLM ([co_rec_ify (StgRec (ids `zip` rhss') :
134 getScBinds (unionLiftInfos rhs_infos))
137 (ids, rhss) = unzip pairs
143 -> LiftM (StgExpr, LiftInfo)
146 liftExpr expr@(StgCon con args lvs) = returnLM (expr, emptyLiftInfo)
147 liftExpr expr@(StgPrim op args lvs) = returnLM (expr, emptyLiftInfo)
149 liftExpr expr@(StgApp (StgLitArg lit) args lvs) = returnLM (expr, emptyLiftInfo)
150 liftExpr expr@(StgApp (StgVarArg v) args lvs)
151 = lookup v `thenLM` \ ~(sc, sc_args) -> -- NB the ~. We don't want to
152 -- poke these bindings too early!
153 returnLM (StgApp (StgVarArg sc) (map StgVarArg sc_args ++ args) lvs,
155 -- The lvs field is probably wrong, but we reconstruct it
156 -- anyway following lambda lifting
158 liftExpr (StgCase scrut lv1 lv2 uniq alts)
159 = liftExpr scrut `thenLM` \ (scrut', scrut_info) ->
160 lift_alts alts `thenLM` \ (alts', alts_info) ->
161 returnLM (StgCase scrut' lv1 lv2 uniq alts', scrut_info `unionLiftInfo` alts_info)
163 lift_alts (StgAlgAlts ty alg_alts deflt)
164 = mapAndUnzipLM lift_alg_alt alg_alts `thenLM` \ (alg_alts', alt_infos) ->
165 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
166 returnLM (StgAlgAlts ty alg_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
168 lift_alts (StgPrimAlts ty prim_alts deflt)
169 = mapAndUnzipLM lift_prim_alt prim_alts `thenLM` \ (prim_alts', alt_infos) ->
170 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
171 returnLM (StgPrimAlts ty prim_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
173 lift_alg_alt (con, args, use_mask, rhs)
174 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
175 returnLM ((con, args, use_mask, rhs'), rhs_info)
177 lift_prim_alt (lit, rhs)
178 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
179 returnLM ((lit, rhs'), rhs_info)
181 lift_deflt StgNoDefault = returnLM (StgNoDefault, emptyLiftInfo)
182 lift_deflt (StgBindDefault var used rhs)
183 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
184 returnLM (StgBindDefault var used rhs', rhs_info)
187 Now the interesting cases. Let no escape isn't lifted. We turn it
188 back into a let, to play safe, because we have to redo that pass after
192 liftExpr (StgLetNoEscape _ _ (StgNonRec binder rhs) body)
193 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
194 liftExpr body `thenLM` \ (body', body_info) ->
195 returnLM (StgLet (StgNonRec binder rhs') body',
196 rhs_info `unionLiftInfo` body_info)
198 liftExpr (StgLetNoEscape _ _ (StgRec pairs) body)
199 = liftExpr body `thenLM` \ (body', body_info) ->
200 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
201 returnLM (StgLet (StgRec (binders `zipEqual` rhss')) body',
202 foldr unionLiftInfo body_info rhs_infos)
204 (binders,rhss) = unzip pairs
208 liftExpr (StgLet (StgNonRec binder rhs) body)
209 | not (isLiftable rhs)
210 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
211 liftExpr body `thenLM` \ (body', body_info) ->
212 returnLM (StgLet (StgNonRec binder rhs') body',
213 rhs_info `unionLiftInfo` body_info)
215 | otherwise -- It's a lambda
216 = -- Do the body of the let
217 fixLM (\ ~(sc_inline, _, _) ->
218 addScInlines [binder] [sc_inline] (
220 ) `thenLM` \ (body', body_info) ->
223 dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
225 -- All occurrences in function position, so lambda lift
226 getFinalFreeVars (rhsFreeVars rhs) `thenLM` \ final_free_vars ->
228 mkScPieces final_free_vars (binder,rhs') `thenLM` \ (sc_inline, sc_bind) ->
232 nonRecScBind rhs_info sc_bind `unionLiftInfo` body_info)
234 ) `thenLM` \ (_, expr', final_info) ->
236 returnLM (expr', final_info)
238 liftExpr (StgLet (StgRec pairs) body)
240 | not (all isLiftableRec rhss)
241 = liftExpr body `thenLM` \ (body', body_info) ->
242 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
243 returnLM (StgLet (StgRec (binders `zipEqual` rhss')) body',
244 foldr unionLiftInfo body_info rhs_infos)
246 | otherwise -- All rhss are liftable
247 = -- Do the body of the let
248 fixLM (\ ~(sc_inlines, _, _) ->
249 addScInlines binders sc_inlines (
251 liftExpr body `thenLM` \ (body', body_info) ->
252 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
254 -- Find the free vars of all the rhss,
255 -- excluding the binders themselves.
256 rhs_free_vars = unionManyIdSets (map rhsFreeVars rhss)
260 rhs_info = unionLiftInfos rhs_infos
262 getFinalFreeVars rhs_free_vars `thenLM` \ final_free_vars ->
264 mapAndUnzipLM (mkScPieces final_free_vars) (binders `zip` rhss')
265 `thenLM` \ (sc_inlines, sc_pairs) ->
266 returnLM (sc_inlines,
268 recScBind rhs_info sc_pairs `unionLiftInfo` body_info)
270 )) `thenLM` \ (_, expr', final_info) ->
272 returnLM (expr', final_info)
274 (binders,rhss) = unzip pairs
278 liftExpr (StgSCC ty cc expr)
279 = liftExpr expr `thenLM` \ (expr2, expr_info) ->
280 returnLM (StgSCC ty cc expr2, expr_info)
283 A binding is liftable if it's a *function* (args not null) and never
284 occurs in an argument position.
287 isLiftable :: StgRhs -> Bool
289 isLiftable (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
291 -- Experimental evidence suggests we should lift only if we will be
292 -- abstracting up to 4 fvs.
294 = if not (null args || -- Not a function
295 unapplied_occ || -- Has an occ with no args at all
296 arg_occ || -- Occurs in arg position
297 length fvs > 4 -- Too many free variables
299 then {-trace ("LL: " ++ show (length fvs))-} True
301 isLiftable other_rhs = False
303 isLiftableRec :: StgRhs -> Bool
305 -- this is just the same as for non-rec, except we only lift to
306 -- abstract up to 1 argument this avoids undoing Static Argument
307 -- Transformation work
309 {- Andre's longer comment about isLiftableRec: 1996/01:
311 A rec binding is "liftable" (according to our heuristics) if:
313 * all occurrences have arguments,
314 * does not occur in an argument position and
315 * has up to *2* free variables (including the rec binding variable
318 The point is: my experiments show that SAT is more important than LL.
319 Therefore if we still want to do LL, for *recursive* functions, we do
320 not want LL to undo what SAT did. We do this by avoiding LL recursive
321 functions that have more than 2 fvs, since if this recursive function
322 was created by SAT (we don't know!), it would have at least 3 fvs: one
323 for the rec binding itself and 2 more for the static arguments (note:
324 this matches with the choice of performing SAT to have at least 2
325 static arguments, if we change things there we should change things
329 isLiftableRec (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
330 = if not (null args || -- Not a function
331 unapplied_occ || -- Has an occ with no args at all
332 arg_occ || -- Occurs in arg position
333 length fvs > 2 -- Too many free variables
335 then {-trace ("LLRec: " ++ show (length fvs))-} True
337 isLiftableRec other_rhs = False
339 rhsFreeVars :: StgRhs -> IdSet
340 rhsFreeVars (StgRhsClosure _ _ fvs _ _ _) = mkIdSet fvs
341 rhsFreeVars other = panic "rhsFreeVars"
344 dontLiftRhs is like liftExpr, except that it does not lift a top-level
345 lambda abstraction. It is used for the right-hand sides of
346 definitions where we've decided *not* to lift: for example, top-level
347 ones or mutually-recursive ones where not all are lambdas.
350 dontLiftRhs :: StgRhs -> LiftM (StgRhs, LiftInfo)
352 dontLiftRhs rhs@(StgRhsCon cc v args) = returnLM (rhs, emptyLiftInfo)
354 dontLiftRhs (StgRhsClosure cc bi fvs upd args body)
355 = liftExpr body `thenLM` \ (body', body_info) ->
356 returnLM (StgRhsClosure cc bi fvs upd args body', body_info)
360 mkScPieces :: IdSet -- Extra args for the supercombinator
361 -> (Id, StgRhs) -- The processed RHS and original Id
362 -> LiftM ((Id,[Id]), -- Replace abstraction with this;
363 -- the set is its free vars
364 (Id,StgRhs)) -- Binding for supercombinator
366 mkScPieces extra_arg_set (id, StgRhsClosure cc bi _ upd args body)
367 = ASSERT( n_args > 0 )
368 -- Construct the rhs of the supercombinator, and its Id
369 newSupercombinator sc_ty arity `thenLM` \ sc_id ->
370 returnLM ((sc_id, extra_args), (sc_id, sc_rhs))
373 extra_args = idSetToList extra_arg_set
374 arity = n_args + length extra_args
376 -- Construct the supercombinator type
377 type_of_original_id = idType id
378 extra_arg_tys = map idType extra_args
379 (tyvars, rest) = splitForAllTy type_of_original_id
380 sc_ty = mkForAllTys tyvars (mkFunTys extra_arg_tys rest)
382 sc_rhs = StgRhsClosure cc bi [] upd (extra_args ++ args) body
386 %************************************************************************
388 \subsection[Lift-monad]{The LiftM monad}
390 %************************************************************************
392 The monad is used only to distribute global stuff, and the unique supply.
395 type LiftM a = LiftFlags
397 -> (IdEnv -- Domain = candidates for lifting
398 (Id, -- The supercombinator
399 [Id]) -- Args to apply it to
404 type LiftFlags = Maybe Int -- No of fvs reqd to float recursive
405 -- binding; Nothing == infinity
408 runLM :: LiftFlags -> UniqSupply -> LiftM a -> a
409 runLM flags us m = m flags us nullIdEnv
411 thenLM :: LiftM a -> (a -> LiftM b) -> LiftM b
412 thenLM m k ci us idenv
413 = k (m ci us1 idenv) ci us2 idenv
415 (us1, us2) = splitUniqSupply us
417 returnLM :: a -> LiftM a
418 returnLM a ci us idenv = a
420 fixLM :: (a -> LiftM a) -> LiftM a
421 fixLM k ci us idenv = r
425 mapLM :: (a -> LiftM b) -> [a] -> LiftM [b]
426 mapLM f [] = returnLM []
427 mapLM f (a:as) = f a `thenLM` \ r ->
428 mapLM f as `thenLM` \ rs ->
431 mapAndUnzipLM :: (a -> LiftM (b,c)) -> [a] -> LiftM ([b],[c])
432 mapAndUnzipLM f [] = returnLM ([],[])
433 mapAndUnzipLM f (a:as) = f a `thenLM` \ (b,c) ->
434 mapAndUnzipLM f as `thenLM` \ (bs,cs) ->
435 returnLM (b:bs, c:cs)
439 newSupercombinator :: Type
443 newSupercombinator ty arity ci us idenv
444 = (mkSysLocal SLIT("sc") uniq ty mkUnknownSrcLoc) -- ToDo: improve location
446 -- ToDo: rm the addIdArity? Just let subsequent stg-saturation pass do it?
450 lookup :: Id -> LiftM (Id,[Id])
452 = case (lookupIdEnv idenv v) of
453 Just result -> result
456 addScInlines :: [Id] -> [(Id,[Id])] -> LiftM a -> LiftM a
457 addScInlines ids values m ci us idenv
460 idenv' = growIdEnvList idenv (ids `zip_lazy` values)
462 -- zip_lazy zips two things together but matches lazily on the
463 -- second argument. This is important, because the ids are know here,
464 -- but the things they are bound to are decided only later
466 zip_lazy (x:xs) ~(y:ys) = (x,y) : zip_lazy xs ys
469 -- The free vars reported by the free-var analyser will include
470 -- some ids, f, which are to be replaced by ($f a b c), where $f
471 -- is the supercombinator. Hence instead of f being a free var,
476 -- f a = ...y1..y2.....
483 -- Here the free vars of g are {f,z}; but f will be lambda-lifted
484 -- with free vars {y1,y2}, so the "real~ free vars of g are {y1,y2,z}.
486 getFinalFreeVars :: IdSet -> LiftM IdSet
488 getFinalFreeVars free_vars ci us idenv
489 = unionManyIdSets (map munge_it (idSetToList free_vars))
491 munge_it :: Id -> IdSet -- Takes a free var and maps it to the "real"
493 munge_it id = case (lookupIdEnv idenv id) of
494 Just (_, args) -> mkIdSet args
495 Nothing -> unitIdSet id
499 %************************************************************************
501 \subsection[Lift-info]{The LiftInfo type}
503 %************************************************************************
506 type LiftInfo = Bag StgBinding -- Float to top
508 emptyLiftInfo = emptyBag
510 unionLiftInfo :: LiftInfo -> LiftInfo -> LiftInfo
511 unionLiftInfo binds1 binds2 = binds1 `unionBags` binds2
513 unionLiftInfos :: [LiftInfo] -> LiftInfo
514 unionLiftInfos infos = foldr unionLiftInfo emptyLiftInfo infos
516 mkScInfo :: StgBinding -> LiftInfo
517 mkScInfo bind = unitBag bind
519 nonRecScBind :: LiftInfo -- From body of supercombinator
520 -> (Id, StgRhs) -- Supercombinator and its rhs
522 nonRecScBind binds (sc_id,sc_rhs) = binds `snocBag` (StgNonRec sc_id sc_rhs)
525 -- In the recursive case, all the SCs from the RHSs of the recursive group
526 -- are dealing with might potentially mention the new, recursive SCs.
527 -- So we flatten the whole lot into a single recursive group.
529 recScBind :: LiftInfo -- From body of supercombinator
530 -> [(Id,StgRhs)] -- Supercombinator rhs
533 recScBind binds pairs = unitBag (co_rec_ify (StgRec pairs : bagToList binds))
535 co_rec_ify :: [StgBinding] -> StgBinding
536 co_rec_ify binds = StgRec (concat (map f binds))
538 f (StgNonRec id rhs) = [(id,rhs)]
539 f (StgRec pairs) = pairs
542 getScBinds :: LiftInfo -> [StgBinding]
543 getScBinds binds = bagToList binds
545 looksLikeSATRhs [(f,StgRhsClosure _ _ _ _ ls _)] (StgApp (StgVarArg f') args _)
546 = (f == f') && (length args == length ls)
547 looksLikeSATRhs _ _ = False