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 (StgVarArg v) args lvs)
152 = lookUp v `thenLM` \ ~(sc, sc_args) -> -- NB the ~. We don't want to
153 -- poke these bindings too early!
154 returnLM (StgApp (StgVarArg sc) (map StgVarArg sc_args ++ args) lvs,
156 -- The lvs field is probably wrong, but we reconstruct it
157 -- anyway following lambda lifting
159 liftExpr (StgCase scrut lv1 lv2 uniq alts)
160 = liftExpr scrut `thenLM` \ (scrut', scrut_info) ->
161 lift_alts alts `thenLM` \ (alts', alts_info) ->
162 returnLM (StgCase scrut' lv1 lv2 uniq alts', scrut_info `unionLiftInfo` alts_info)
164 lift_alts (StgAlgAlts ty alg_alts deflt)
165 = mapAndUnzipLM lift_alg_alt alg_alts `thenLM` \ (alg_alts', alt_infos) ->
166 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
167 returnLM (StgAlgAlts ty alg_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
169 lift_alts (StgPrimAlts ty prim_alts deflt)
170 = mapAndUnzipLM lift_prim_alt prim_alts `thenLM` \ (prim_alts', alt_infos) ->
171 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
172 returnLM (StgPrimAlts ty prim_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
174 lift_alg_alt (con, args, use_mask, rhs)
175 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
176 returnLM ((con, args, use_mask, rhs'), rhs_info)
178 lift_prim_alt (lit, rhs)
179 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
180 returnLM ((lit, rhs'), rhs_info)
182 lift_deflt StgNoDefault = returnLM (StgNoDefault, emptyLiftInfo)
183 lift_deflt (StgBindDefault var used rhs)
184 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
185 returnLM (StgBindDefault var used rhs', rhs_info)
188 Now the interesting cases. Let no escape isn't lifted. We turn it
189 back into a let, to play safe, because we have to redo that pass after
193 liftExpr (StgLetNoEscape _ _ (StgNonRec binder rhs) body)
194 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
195 liftExpr body `thenLM` \ (body', body_info) ->
196 returnLM (StgLet (StgNonRec binder rhs') body',
197 rhs_info `unionLiftInfo` body_info)
199 liftExpr (StgLetNoEscape _ _ (StgRec pairs) body)
200 = liftExpr body `thenLM` \ (body', body_info) ->
201 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
202 returnLM (StgLet (StgRec (zipEqual "liftExpr" binders rhss')) body',
203 foldr unionLiftInfo body_info rhs_infos)
205 (binders,rhss) = unzip pairs
209 liftExpr (StgLet (StgNonRec binder rhs) body)
210 | not (isLiftable rhs)
211 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
212 liftExpr body `thenLM` \ (body', body_info) ->
213 returnLM (StgLet (StgNonRec binder rhs') body',
214 rhs_info `unionLiftInfo` body_info)
216 | otherwise -- It's a lambda
217 = -- Do the body of the let
218 fixLM (\ ~(sc_inline, _, _) ->
219 addScInlines [binder] [sc_inline] (
221 ) `thenLM` \ (body', body_info) ->
224 dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
226 -- All occurrences in function position, so lambda lift
227 getFinalFreeVars (rhsFreeVars rhs) `thenLM` \ final_free_vars ->
229 mkScPieces final_free_vars (binder,rhs') `thenLM` \ (sc_inline, sc_bind) ->
233 nonRecScBind rhs_info sc_bind `unionLiftInfo` body_info)
235 ) `thenLM` \ (_, expr', final_info) ->
237 returnLM (expr', final_info)
239 liftExpr (StgLet (StgRec pairs) body)
241 | not (all isLiftableRec rhss)
242 = liftExpr body `thenLM` \ (body', body_info) ->
243 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
244 returnLM (StgLet (StgRec (zipEqual "liftExpr2" binders rhss')) body',
245 foldr unionLiftInfo body_info rhs_infos)
247 | otherwise -- All rhss are liftable
248 = -- Do the body of the let
249 fixLM (\ ~(sc_inlines, _, _) ->
250 addScInlines binders sc_inlines (
252 liftExpr body `thenLM` \ (body', body_info) ->
253 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
255 -- Find the free vars of all the rhss,
256 -- excluding the binders themselves.
257 rhs_free_vars = unionManyIdSets (map rhsFreeVars rhss)
261 rhs_info = unionLiftInfos rhs_infos
263 getFinalFreeVars rhs_free_vars `thenLM` \ final_free_vars ->
265 mapAndUnzipLM (mkScPieces final_free_vars) (binders `zip` rhss')
266 `thenLM` \ (sc_inlines, sc_pairs) ->
267 returnLM (sc_inlines,
269 recScBind rhs_info sc_pairs `unionLiftInfo` body_info)
271 )) `thenLM` \ (_, expr', final_info) ->
273 returnLM (expr', final_info)
275 (binders,rhss) = unzip pairs
279 liftExpr (StgSCC ty cc expr)
280 = liftExpr expr `thenLM` \ (expr2, expr_info) ->
281 returnLM (StgSCC ty cc expr2, expr_info)
284 A binding is liftable if it's a *function* (args not null) and never
285 occurs in an argument position.
288 isLiftable :: StgRhs -> Bool
290 isLiftable (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
292 -- Experimental evidence suggests we should lift only if we will be
293 -- abstracting up to 4 fvs.
295 = if not (null args || -- Not a function
296 unapplied_occ || -- Has an occ with no args at all
297 arg_occ || -- Occurs in arg position
298 length fvs > 4 -- Too many free variables
300 then {-trace ("LL: " ++ show (length fvs))-} True
302 isLiftable other_rhs = False
304 isLiftableRec :: StgRhs -> Bool
306 -- this is just the same as for non-rec, except we only lift to
307 -- abstract up to 1 argument this avoids undoing Static Argument
308 -- Transformation work
310 {- Andre's longer comment about isLiftableRec: 1996/01:
312 A rec binding is "liftable" (according to our heuristics) if:
314 * all occurrences have arguments,
315 * does not occur in an argument position and
316 * has up to *2* free variables (including the rec binding variable
319 The point is: my experiments show that SAT is more important than LL.
320 Therefore if we still want to do LL, for *recursive* functions, we do
321 not want LL to undo what SAT did. We do this by avoiding LL recursive
322 functions that have more than 2 fvs, since if this recursive function
323 was created by SAT (we don't know!), it would have at least 3 fvs: one
324 for the rec binding itself and 2 more for the static arguments (note:
325 this matches with the choice of performing SAT to have at least 2
326 static arguments, if we change things there we should change things
330 isLiftableRec (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
331 = if not (null args || -- Not a function
332 unapplied_occ || -- Has an occ with no args at all
333 arg_occ || -- Occurs in arg position
334 length fvs > 2 -- Too many free variables
336 then {-trace ("LLRec: " ++ show (length fvs))-} True
338 isLiftableRec other_rhs = False
340 rhsFreeVars :: StgRhs -> IdSet
341 rhsFreeVars (StgRhsClosure _ _ fvs _ _ _) = mkIdSet fvs
342 rhsFreeVars other = panic "rhsFreeVars"
345 dontLiftRhs is like liftExpr, except that it does not lift a top-level
346 lambda abstraction. It is used for the right-hand sides of
347 definitions where we've decided *not* to lift: for example, top-level
348 ones or mutually-recursive ones where not all are lambdas.
351 dontLiftRhs :: StgRhs -> LiftM (StgRhs, LiftInfo)
353 dontLiftRhs rhs@(StgRhsCon cc v args) = returnLM (rhs, emptyLiftInfo)
355 dontLiftRhs (StgRhsClosure cc bi fvs upd args body)
356 = liftExpr body `thenLM` \ (body', body_info) ->
357 returnLM (StgRhsClosure cc bi fvs upd args body', body_info)
361 mkScPieces :: IdSet -- Extra args for the supercombinator
362 -> (Id, StgRhs) -- The processed RHS and original Id
363 -> LiftM ((Id,[Id]), -- Replace abstraction with this;
364 -- the set is its free vars
365 (Id,StgRhs)) -- Binding for supercombinator
367 mkScPieces extra_arg_set (id, StgRhsClosure cc bi _ upd args body)
368 = ASSERT( n_args > 0 )
369 -- Construct the rhs of the supercombinator, and its Id
370 newSupercombinator sc_ty arity `thenLM` \ sc_id ->
371 returnLM ((sc_id, extra_args), (sc_id, sc_rhs))
374 extra_args = idSetToList extra_arg_set
375 arity = n_args + length extra_args
377 -- Construct the supercombinator type
378 type_of_original_id = idType id
379 extra_arg_tys = map idType extra_args
380 (tyvars, rest) = splitForAllTy type_of_original_id
381 sc_ty = mkForAllTys tyvars (mkFunTys extra_arg_tys rest)
383 sc_rhs = StgRhsClosure cc bi [] upd (extra_args ++ args) body
387 %************************************************************************
389 \subsection[Lift-monad]{The LiftM monad}
391 %************************************************************************
393 The monad is used only to distribute global stuff, and the unique supply.
396 type LiftM a = LiftFlags
398 -> (IdEnv -- Domain = candidates for lifting
399 (Id, -- The supercombinator
400 [Id]) -- Args to apply it to
405 type LiftFlags = Maybe Int -- No of fvs reqd to float recursive
406 -- binding; Nothing == infinity
409 runLM :: LiftFlags -> UniqSupply -> LiftM a -> a
410 runLM flags us m = m flags us nullIdEnv
412 thenLM :: LiftM a -> (a -> LiftM b) -> LiftM b
413 thenLM m k ci us idenv
414 = k (m ci us1 idenv) ci us2 idenv
416 (us1, us2) = splitUniqSupply us
418 returnLM :: a -> LiftM a
419 returnLM a ci us idenv = a
421 fixLM :: (a -> LiftM a) -> LiftM a
422 fixLM k ci us idenv = r
426 mapLM :: (a -> LiftM b) -> [a] -> LiftM [b]
427 mapLM f [] = returnLM []
428 mapLM f (a:as) = f a `thenLM` \ r ->
429 mapLM f as `thenLM` \ rs ->
432 mapAndUnzipLM :: (a -> LiftM (b,c)) -> [a] -> LiftM ([b],[c])
433 mapAndUnzipLM f [] = returnLM ([],[])
434 mapAndUnzipLM f (a:as) = f a `thenLM` \ (b,c) ->
435 mapAndUnzipLM f as `thenLM` \ (bs,cs) ->
436 returnLM (b:bs, c:cs)
440 newSupercombinator :: Type
444 newSupercombinator ty arity ci us idenv
445 = (mkSysLocal SLIT("sc") uniq ty noSrcLoc) -- ToDo: improve location
446 `addIdArity` exactArity arity
447 -- ToDo: rm the addIdArity? Just let subsequent stg-saturation pass do it?
451 lookUp :: Id -> LiftM (Id,[Id])
453 = case (lookupIdEnv idenv v) of
454 Just result -> result
457 addScInlines :: [Id] -> [(Id,[Id])] -> LiftM a -> LiftM a
458 addScInlines ids values m ci us idenv
461 idenv' = growIdEnvList idenv (ids `zip_lazy` values)
463 -- zip_lazy zips two things together but matches lazily on the
464 -- second argument. This is important, because the ids are know here,
465 -- but the things they are bound to are decided only later
467 zip_lazy (x:xs) ~(y:ys) = (x,y) : zip_lazy xs ys
470 -- The free vars reported by the free-var analyser will include
471 -- some ids, f, which are to be replaced by ($f a b c), where $f
472 -- is the supercombinator. Hence instead of f being a free var,
477 -- f a = ...y1..y2.....
484 -- Here the free vars of g are {f,z}; but f will be lambda-lifted
485 -- with free vars {y1,y2}, so the "real~ free vars of g are {y1,y2,z}.
487 getFinalFreeVars :: IdSet -> LiftM IdSet
489 getFinalFreeVars free_vars ci us idenv
490 = unionManyIdSets (map munge_it (idSetToList free_vars))
492 munge_it :: Id -> IdSet -- Takes a free var and maps it to the "real"
494 munge_it id = case (lookupIdEnv idenv id) of
495 Just (_, args) -> mkIdSet args
496 Nothing -> unitIdSet id
500 %************************************************************************
502 \subsection[Lift-info]{The LiftInfo type}
504 %************************************************************************
507 type LiftInfo = Bag StgBinding -- Float to top
509 emptyLiftInfo = emptyBag
511 unionLiftInfo :: LiftInfo -> LiftInfo -> LiftInfo
512 unionLiftInfo binds1 binds2 = binds1 `unionBags` binds2
514 unionLiftInfos :: [LiftInfo] -> LiftInfo
515 unionLiftInfos infos = foldr unionLiftInfo emptyLiftInfo infos
517 mkScInfo :: StgBinding -> LiftInfo
518 mkScInfo bind = unitBag bind
520 nonRecScBind :: LiftInfo -- From body of supercombinator
521 -> (Id, StgRhs) -- Supercombinator and its rhs
523 nonRecScBind binds (sc_id,sc_rhs) = binds `snocBag` (StgNonRec sc_id sc_rhs)
526 -- In the recursive case, all the SCs from the RHSs of the recursive group
527 -- are dealing with might potentially mention the new, recursive SCs.
528 -- So we flatten the whole lot into a single recursive group.
530 recScBind :: LiftInfo -- From body of supercombinator
531 -> [(Id,StgRhs)] -- Supercombinator rhs
534 recScBind binds pairs = unitBag (co_rec_ify (StgRec pairs : bagToList binds))
536 co_rec_ify :: [StgBinding] -> StgBinding
537 co_rec_ify binds = StgRec (concat (map f binds))
539 f (StgNonRec id rhs) = [(id,rhs)]
540 f (StgRec pairs) = pairs
543 getScBinds :: LiftInfo -> [StgBinding]
544 getScBinds binds = bagToList binds
546 looksLikeSATRhs [(f,StgRhsClosure _ _ _ _ ls _)] (StgApp (StgVarArg f') args _)
547 = (f == f') && (length args == length ls)
548 looksLikeSATRhs _ _ = False