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, SYN_IE(IdSet),
19 nullIdEnv, growIdEnvList, lookupIdEnv, SYN_IE(IdEnv)
21 import IdInfo ( ArityInfo, exactArity )
22 import SrcLoc ( noSrcLoc )
23 import Type ( splitForAllTy, mkForAllTys, mkFunTys )
24 import UniqSupply ( getUnique, splitUniqSupply )
25 import Util ( zipEqual, panic, assertPanic )
28 This is the lambda lifter. It turns lambda abstractions into
29 supercombinators on a selective basis:
31 * Let-no-escaped bindings are never lifted. That's one major reason
32 why the lambda lifter is done in STG.
34 * Non-recursive bindings whose RHS is a lambda abstractions are lifted,
35 provided all the occurrences of the bound variable is in a function
36 postition. In this example, f will be lifted:
44 $f p q r x = e -- Supercombinator
46 ..($f p q r a1)...($f p q r a2)...
48 NOTE that the original binding is eliminated.
50 But in this case, f won't be lifted:
57 Why? Because we have to heap-allocate a closure for f thus:
59 $f p q r x = e -- Supercombinator
64 ..(g f)...($f p q r a2)..
66 so it might as well be the original lambda abstraction.
68 We also do not lift if the function has an occurrence with no arguments, e.g.
74 as this form is more efficient than if we create a partial application
76 $f p q r x = e -- Supercombinator
80 * Recursive bindings *all* of whose RHSs are lambda abstractions are
82 - all the occurrences of all the binders are in a function position
83 - there aren't ``too many'' free variables.
85 Same reasoning as before for the function-position stuff. The ``too many
86 free variable'' part comes from considering the (potentially many)
87 recursive calls, which may now have lots of free vars.
90 * 2 might be already ``too many'' variables to abstract.
91 The problem is that the increase in the number of free variables
92 of closures refering to the lifted function (which is always # of
93 abstracted args - 1) may increase heap allocation a lot.
94 Expeiments are being done to check this...
95 * We do not lambda lift if the function has at least one occurrence
96 without any arguments. This caused lots of problems. Ex:
97 h = \ x -> ... let y = ...
98 in let let f = \x -> ...y...
102 h = \ x -> ... let y = ...
105 now f y is a partial application, so it will be updated, and this
109 --- NOT RELEVANT FOR STG ----
110 * All ``lone'' lambda abstractions are lifted. Notably this means lambda
112 - in a case alternative: case e of True -> (\x->b)
113 - in the body of a let: let x=e in (\y->b)
114 -----------------------------
116 %************************************************************************
118 \subsection[Lift-expressions]{The main function: liftExpr}
120 %************************************************************************
123 liftProgram :: UniqSupply -> [StgBinding] -> [StgBinding]
124 liftProgram us prog = concat (runLM Nothing us (mapLM liftTopBind prog))
127 liftTopBind :: StgBinding -> LiftM [StgBinding]
128 liftTopBind (StgNonRec id rhs)
129 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
130 returnLM (getScBinds rhs_info ++ [StgNonRec id rhs'])
132 liftTopBind (StgRec pairs)
133 = mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
134 returnLM ([co_rec_ify (StgRec (ids `zip` rhss') :
135 getScBinds (unionLiftInfos rhs_infos))
138 (ids, rhss) = unzip pairs
144 -> LiftM (StgExpr, LiftInfo)
147 liftExpr expr@(StgCon con args lvs) = returnLM (expr, emptyLiftInfo)
148 liftExpr expr@(StgPrim op args lvs) = returnLM (expr, emptyLiftInfo)
150 liftExpr expr@(StgApp (StgLitArg lit) args lvs) = returnLM (expr, emptyLiftInfo)
151 liftExpr expr@(StgApp (StgConArg con) args lvs) = returnLM (expr, emptyLiftInfo)
152 liftExpr expr@(StgApp (StgVarArg v) args lvs)
153 = lookUp v `thenLM` \ ~(sc, sc_args) -> -- NB the ~. We don't want to
154 -- poke these bindings too early!
155 returnLM (StgApp (StgVarArg sc) (map StgVarArg sc_args ++ args) lvs,
157 -- The lvs field is probably wrong, but we reconstruct it
158 -- anyway following lambda lifting
160 liftExpr (StgCase scrut lv1 lv2 uniq alts)
161 = liftExpr scrut `thenLM` \ (scrut', scrut_info) ->
162 lift_alts alts `thenLM` \ (alts', alts_info) ->
163 returnLM (StgCase scrut' lv1 lv2 uniq alts', scrut_info `unionLiftInfo` alts_info)
165 lift_alts (StgAlgAlts ty alg_alts deflt)
166 = mapAndUnzipLM lift_alg_alt alg_alts `thenLM` \ (alg_alts', alt_infos) ->
167 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
168 returnLM (StgAlgAlts ty alg_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
170 lift_alts (StgPrimAlts ty prim_alts deflt)
171 = mapAndUnzipLM lift_prim_alt prim_alts `thenLM` \ (prim_alts', alt_infos) ->
172 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
173 returnLM (StgPrimAlts ty prim_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
175 lift_alg_alt (con, args, use_mask, rhs)
176 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
177 returnLM ((con, args, use_mask, rhs'), rhs_info)
179 lift_prim_alt (lit, rhs)
180 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
181 returnLM ((lit, rhs'), rhs_info)
183 lift_deflt StgNoDefault = returnLM (StgNoDefault, emptyLiftInfo)
184 lift_deflt (StgBindDefault var used rhs)
185 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
186 returnLM (StgBindDefault var used rhs', rhs_info)
189 Now the interesting cases. Let no escape isn't lifted. We turn it
190 back into a let, to play safe, because we have to redo that pass after
194 liftExpr (StgLetNoEscape _ _ (StgNonRec binder rhs) body)
195 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
196 liftExpr body `thenLM` \ (body', body_info) ->
197 returnLM (StgLet (StgNonRec binder rhs') body',
198 rhs_info `unionLiftInfo` body_info)
200 liftExpr (StgLetNoEscape _ _ (StgRec pairs) body)
201 = liftExpr body `thenLM` \ (body', body_info) ->
202 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
203 returnLM (StgLet (StgRec (zipEqual "liftExpr" binders rhss')) body',
204 foldr unionLiftInfo body_info rhs_infos)
206 (binders,rhss) = unzip pairs
210 liftExpr (StgLet (StgNonRec binder rhs) body)
211 | not (isLiftable rhs)
212 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
213 liftExpr body `thenLM` \ (body', body_info) ->
214 returnLM (StgLet (StgNonRec binder rhs') body',
215 rhs_info `unionLiftInfo` body_info)
217 | otherwise -- It's a lambda
218 = -- Do the body of the let
219 fixLM (\ ~(sc_inline, _, _) ->
220 addScInlines [binder] [sc_inline] (
222 ) `thenLM` \ (body', body_info) ->
225 dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
227 -- All occurrences in function position, so lambda lift
228 getFinalFreeVars (rhsFreeVars rhs) `thenLM` \ final_free_vars ->
230 mkScPieces final_free_vars (binder,rhs') `thenLM` \ (sc_inline, sc_bind) ->
234 nonRecScBind rhs_info sc_bind `unionLiftInfo` body_info)
236 ) `thenLM` \ (_, expr', final_info) ->
238 returnLM (expr', final_info)
240 liftExpr (StgLet (StgRec pairs) body)
242 | not (all isLiftableRec rhss)
243 = liftExpr body `thenLM` \ (body', body_info) ->
244 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
245 returnLM (StgLet (StgRec (zipEqual "liftExpr2" binders rhss')) body',
246 foldr unionLiftInfo body_info rhs_infos)
248 | otherwise -- All rhss are liftable
249 = -- Do the body of the let
250 fixLM (\ ~(sc_inlines, _, _) ->
251 addScInlines binders sc_inlines (
253 liftExpr body `thenLM` \ (body', body_info) ->
254 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
256 -- Find the free vars of all the rhss,
257 -- excluding the binders themselves.
258 rhs_free_vars = unionManyIdSets (map rhsFreeVars rhss)
262 rhs_info = unionLiftInfos rhs_infos
264 getFinalFreeVars rhs_free_vars `thenLM` \ final_free_vars ->
266 mapAndUnzipLM (mkScPieces final_free_vars) (binders `zip` rhss')
267 `thenLM` \ (sc_inlines, sc_pairs) ->
268 returnLM (sc_inlines,
270 recScBind rhs_info sc_pairs `unionLiftInfo` body_info)
272 )) `thenLM` \ (_, expr', final_info) ->
274 returnLM (expr', final_info)
276 (binders,rhss) = unzip pairs
280 liftExpr (StgSCC ty cc expr)
281 = liftExpr expr `thenLM` \ (expr2, expr_info) ->
282 returnLM (StgSCC ty cc expr2, expr_info)
285 A binding is liftable if it's a *function* (args not null) and never
286 occurs in an argument position.
289 isLiftable :: StgRhs -> Bool
291 isLiftable (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
293 -- Experimental evidence suggests we should lift only if we will be
294 -- abstracting up to 4 fvs.
296 = if not (null args || -- Not a function
297 unapplied_occ || -- Has an occ with no args at all
298 arg_occ || -- Occurs in arg position
299 length fvs > 4 -- Too many free variables
301 then {-trace ("LL: " ++ show (length fvs))-} True
303 isLiftable other_rhs = False
305 isLiftableRec :: StgRhs -> Bool
307 -- this is just the same as for non-rec, except we only lift to
308 -- abstract up to 1 argument this avoids undoing Static Argument
309 -- Transformation work
311 {- Andre's longer comment about isLiftableRec: 1996/01:
313 A rec binding is "liftable" (according to our heuristics) if:
315 * all occurrences have arguments,
316 * does not occur in an argument position and
317 * has up to *2* free variables (including the rec binding variable
320 The point is: my experiments show that SAT is more important than LL.
321 Therefore if we still want to do LL, for *recursive* functions, we do
322 not want LL to undo what SAT did. We do this by avoiding LL recursive
323 functions that have more than 2 fvs, since if this recursive function
324 was created by SAT (we don't know!), it would have at least 3 fvs: one
325 for the rec binding itself and 2 more for the static arguments (note:
326 this matches with the choice of performing SAT to have at least 2
327 static arguments, if we change things there we should change things
331 isLiftableRec (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
332 = if not (null args || -- Not a function
333 unapplied_occ || -- Has an occ with no args at all
334 arg_occ || -- Occurs in arg position
335 length fvs > 2 -- Too many free variables
337 then {-trace ("LLRec: " ++ show (length fvs))-} True
339 isLiftableRec other_rhs = False
341 rhsFreeVars :: StgRhs -> IdSet
342 rhsFreeVars (StgRhsClosure _ _ fvs _ _ _) = mkIdSet fvs
343 rhsFreeVars other = panic "rhsFreeVars"
346 dontLiftRhs is like liftExpr, except that it does not lift a top-level
347 lambda abstraction. It is used for the right-hand sides of
348 definitions where we've decided *not* to lift: for example, top-level
349 ones or mutually-recursive ones where not all are lambdas.
352 dontLiftRhs :: StgRhs -> LiftM (StgRhs, LiftInfo)
354 dontLiftRhs rhs@(StgRhsCon cc v args) = returnLM (rhs, emptyLiftInfo)
356 dontLiftRhs (StgRhsClosure cc bi fvs upd args body)
357 = liftExpr body `thenLM` \ (body', body_info) ->
358 returnLM (StgRhsClosure cc bi fvs upd args body', body_info)
362 mkScPieces :: IdSet -- Extra args for the supercombinator
363 -> (Id, StgRhs) -- The processed RHS and original Id
364 -> LiftM ((Id,[Id]), -- Replace abstraction with this;
365 -- the set is its free vars
366 (Id,StgRhs)) -- Binding for supercombinator
368 mkScPieces extra_arg_set (id, StgRhsClosure cc bi _ upd args body)
369 = ASSERT( n_args > 0 )
370 -- Construct the rhs of the supercombinator, and its Id
371 newSupercombinator sc_ty arity `thenLM` \ sc_id ->
372 returnLM ((sc_id, extra_args), (sc_id, sc_rhs))
375 extra_args = idSetToList extra_arg_set
376 arity = n_args + length extra_args
378 -- Construct the supercombinator type
379 type_of_original_id = idType id
380 extra_arg_tys = map idType extra_args
381 (tyvars, rest) = splitForAllTy type_of_original_id
382 sc_ty = mkForAllTys tyvars (mkFunTys extra_arg_tys rest)
384 sc_rhs = StgRhsClosure cc bi [] upd (extra_args ++ args) body
388 %************************************************************************
390 \subsection[Lift-monad]{The LiftM monad}
392 %************************************************************************
394 The monad is used only to distribute global stuff, and the unique supply.
397 type LiftM a = LiftFlags
399 -> (IdEnv -- Domain = candidates for lifting
400 (Id, -- The supercombinator
401 [Id]) -- Args to apply it to
406 type LiftFlags = Maybe Int -- No of fvs reqd to float recursive
407 -- binding; Nothing == infinity
410 runLM :: LiftFlags -> UniqSupply -> LiftM a -> a
411 runLM flags us m = m flags us nullIdEnv
413 thenLM :: LiftM a -> (a -> LiftM b) -> LiftM b
414 thenLM m k ci us idenv
415 = k (m ci us1 idenv) ci us2 idenv
417 (us1, us2) = splitUniqSupply us
419 returnLM :: a -> LiftM a
420 returnLM a ci us idenv = a
422 fixLM :: (a -> LiftM a) -> LiftM a
423 fixLM k ci us idenv = r
427 mapLM :: (a -> LiftM b) -> [a] -> LiftM [b]
428 mapLM f [] = returnLM []
429 mapLM f (a:as) = f a `thenLM` \ r ->
430 mapLM f as `thenLM` \ rs ->
433 mapAndUnzipLM :: (a -> LiftM (b,c)) -> [a] -> LiftM ([b],[c])
434 mapAndUnzipLM f [] = returnLM ([],[])
435 mapAndUnzipLM f (a:as) = f a `thenLM` \ (b,c) ->
436 mapAndUnzipLM f as `thenLM` \ (bs,cs) ->
437 returnLM (b:bs, c:cs)
441 newSupercombinator :: Type
445 newSupercombinator ty arity ci us idenv
446 = (mkSysLocal SLIT("sc") uniq ty noSrcLoc) -- ToDo: improve location
447 `addIdArity` exactArity arity
448 -- ToDo: rm the addIdArity? Just let subsequent stg-saturation pass do it?
452 lookUp :: Id -> LiftM (Id,[Id])
454 = case (lookupIdEnv idenv v) of
455 Just result -> result
458 addScInlines :: [Id] -> [(Id,[Id])] -> LiftM a -> LiftM a
459 addScInlines ids values m ci us idenv
462 idenv' = growIdEnvList idenv (ids `zip_lazy` values)
464 -- zip_lazy zips two things together but matches lazily on the
465 -- second argument. This is important, because the ids are know here,
466 -- but the things they are bound to are decided only later
468 zip_lazy (x:xs) ~(y:ys) = (x,y) : zip_lazy xs ys
471 -- The free vars reported by the free-var analyser will include
472 -- some ids, f, which are to be replaced by ($f a b c), where $f
473 -- is the supercombinator. Hence instead of f being a free var,
478 -- f a = ...y1..y2.....
485 -- Here the free vars of g are {f,z}; but f will be lambda-lifted
486 -- with free vars {y1,y2}, so the "real~ free vars of g are {y1,y2,z}.
488 getFinalFreeVars :: IdSet -> LiftM IdSet
490 getFinalFreeVars free_vars ci us idenv
491 = unionManyIdSets (map munge_it (idSetToList free_vars))
493 munge_it :: Id -> IdSet -- Takes a free var and maps it to the "real"
495 munge_it id = case (lookupIdEnv idenv id) of
496 Just (_, args) -> mkIdSet args
497 Nothing -> unitIdSet id
501 %************************************************************************
503 \subsection[Lift-info]{The LiftInfo type}
505 %************************************************************************
508 type LiftInfo = Bag StgBinding -- Float to top
510 emptyLiftInfo = emptyBag
512 unionLiftInfo :: LiftInfo -> LiftInfo -> LiftInfo
513 unionLiftInfo binds1 binds2 = binds1 `unionBags` binds2
515 unionLiftInfos :: [LiftInfo] -> LiftInfo
516 unionLiftInfos infos = foldr unionLiftInfo emptyLiftInfo infos
518 mkScInfo :: StgBinding -> LiftInfo
519 mkScInfo bind = unitBag bind
521 nonRecScBind :: LiftInfo -- From body of supercombinator
522 -> (Id, StgRhs) -- Supercombinator and its rhs
524 nonRecScBind binds (sc_id,sc_rhs) = binds `snocBag` (StgNonRec sc_id sc_rhs)
527 -- In the recursive case, all the SCs from the RHSs of the recursive group
528 -- are dealing with might potentially mention the new, recursive SCs.
529 -- So we flatten the whole lot into a single recursive group.
531 recScBind :: LiftInfo -- From body of supercombinator
532 -> [(Id,StgRhs)] -- Supercombinator rhs
535 recScBind binds pairs = unitBag (co_rec_ify (StgRec pairs : bagToList binds))
537 co_rec_ify :: [StgBinding] -> StgBinding
538 co_rec_ify binds = StgRec (concat (map f binds))
540 f (StgNonRec id rhs) = [(id,rhs)]
541 f (StgRec pairs) = pairs
544 getScBinds :: LiftInfo -> [StgBinding]
545 getScBinds binds = bagToList binds
547 looksLikeSATRhs [(f,StgRhsClosure _ _ _ _ ls _)] (StgApp (StgVarArg f') args _)
548 = (f == f') && (length args == length ls)
549 looksLikeSATRhs _ _ = False