2 % (c) The AQUA Project, Glasgow University, 1994-1995
4 \section[LambdaLift]{A STG-code lambda lifter}
7 #include "HsVersions.h"
9 module LambdaLift ( liftProgram ) where
13 import AbsUniType ( mkForallTy, splitForalls, glueTyArgs,
14 UniType, RhoType(..), TauType(..)
17 import Id ( mkSysLocal, getIdUniType, addIdArity, Id )
21 import SrcLoc ( mkUnknownSrcLoc, SrcLoc )
26 This is the lambda lifter. It turns lambda abstractions into
27 supercombinators on a selective basis:
29 * Let-no-escaped bindings are never lifted. That's one major reason
30 why the lambda lifter is done in STG.
32 * Non-recursive bindings whose RHS is a lambda abstractions are lifted,
33 provided all the occurrences of the bound variable is in a function
34 postition. In this example, f will be lifted:
42 $f p q r x = e -- Supercombinator
44 ..($f p q r a1)...($f p q r a2)...
46 NOTE that the original binding is eliminated.
48 But in this case, f won't be lifted:
55 Why? Because we have to heap-allocate a closure for f thus:
57 $f p q r x = e -- Supercombinator
62 ..(g f)...($f p q r a2)..
64 so it might as well be the original lambda abstraction.
66 We also do not lift if the function has an occurrence with no arguments, e.g.
72 as this form is more efficient than if we create a partial application
74 $f p q r x = e -- Supercombinator
78 * Recursive bindings *all* of whose RHSs are lambda abstractions are
80 - all the occurrences of all the binders are in a function position
81 - there aren't ``too many'' free variables.
83 Same reasoning as before for the function-position stuff. The ``too many
84 free variable'' part comes from considering the (potentially many)
85 recursive calls, which may now have lots of free vars.
88 * 2 might be already ``too many'' variables to abstract.
89 The problem is that the increase in the number of free variables
90 of closures refering to the lifted function (which is always # of
91 abstracted args - 1) may increase heap allocation a lot.
92 Expeiments are being done to check this...
93 * We do not lambda lift if the function has at least one occurrence
94 without any arguments. This caused lots of problems. Ex:
95 h = \ x -> ... let y = ...
96 in let let f = \x -> ...y...
100 h = \ x -> ... let y = ...
103 now f y is a partial application, so it will be updated, and this
107 --- NOT RELEVANT FOR STG ----
108 * All ``lone'' lambda abstractions are lifted. Notably this means lambda
110 - in a case alternative: case e of True -> (\x->b)
111 - in the body of a let: let x=e in (\y->b)
112 -----------------------------
114 %************************************************************************
116 \subsection[Lift-expressions]{The main function: liftExpr}
118 %************************************************************************
121 liftProgram :: SplitUniqSupply -> [PlainStgBinding] -> [PlainStgBinding]
122 liftProgram us prog = concat (runLM Nothing us (mapLM liftTopBind prog))
125 liftTopBind :: PlainStgBinding -> LiftM [PlainStgBinding]
126 liftTopBind (StgNonRec id rhs)
127 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
128 returnLM (getScBinds rhs_info ++ [StgNonRec id rhs'])
130 liftTopBind (StgRec pairs)
131 = mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
132 returnLM ([co_rec_ify (StgRec (ids `zip` rhss') :
133 getScBinds (unionLiftInfos rhs_infos))
136 (ids, rhss) = unzip pairs
141 liftExpr :: PlainStgExpr
142 -> LiftM (PlainStgExpr, LiftInfo)
145 liftExpr expr@(StgConApp con args lvs) = returnLM (expr, emptyLiftInfo)
146 liftExpr expr@(StgPrimApp op args lvs) = returnLM (expr, emptyLiftInfo)
148 liftExpr expr@(StgApp (StgLitAtom lit) args lvs) = returnLM (expr, emptyLiftInfo)
149 liftExpr expr@(StgApp (StgVarAtom v) args lvs)
150 = lookup v `thenLM` \ ~(sc, sc_args) -> -- NB the ~. We don't want to
151 -- poke these bindings too early!
152 returnLM (StgApp (StgVarAtom sc) (map StgVarAtom sc_args ++ args) lvs,
154 -- The lvs field is probably wrong, but we reconstruct it
155 -- anyway following lambda lifting
157 liftExpr (StgCase scrut lv1 lv2 uniq alts)
158 = liftExpr scrut `thenLM` \ (scrut', scrut_info) ->
159 lift_alts alts `thenLM` \ (alts', alts_info) ->
160 returnLM (StgCase scrut' lv1 lv2 uniq alts', scrut_info `unionLiftInfo` alts_info)
162 lift_alts (StgAlgAlts ty alg_alts deflt)
163 = mapAndUnzipLM lift_alg_alt alg_alts `thenLM` \ (alg_alts', alt_infos) ->
164 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
165 returnLM (StgAlgAlts ty alg_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
167 lift_alts (StgPrimAlts ty prim_alts deflt)
168 = mapAndUnzipLM lift_prim_alt prim_alts `thenLM` \ (prim_alts', alt_infos) ->
169 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
170 returnLM (StgPrimAlts ty prim_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
172 lift_alg_alt (con, args, use_mask, rhs)
173 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
174 returnLM ((con, args, use_mask, rhs'), rhs_info)
176 lift_prim_alt (lit, rhs)
177 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
178 returnLM ((lit, rhs'), rhs_info)
180 lift_deflt StgNoDefault = returnLM (StgNoDefault, emptyLiftInfo)
181 lift_deflt (StgBindDefault var used rhs)
182 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
183 returnLM (StgBindDefault var used rhs', rhs_info)
186 Now the interesting cases. Let no escape isn't lifted. We turn it
187 back into a let, to play safe, because we have to redo that pass after
191 liftExpr (StgLetNoEscape _ _ (StgNonRec binder rhs) body)
192 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
193 liftExpr body `thenLM` \ (body', body_info) ->
194 returnLM (StgLet (StgNonRec binder rhs') body',
195 rhs_info `unionLiftInfo` body_info)
197 liftExpr (StgLetNoEscape _ _ (StgRec pairs) body)
198 = liftExpr body `thenLM` \ (body', body_info) ->
199 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
200 returnLM (StgLet (StgRec (binders `zipEqual` rhss')) body',
201 foldr unionLiftInfo body_info rhs_infos)
203 (binders,rhss) = unzip pairs
207 liftExpr (StgLet (StgNonRec binder rhs) body)
208 | not (isLiftable rhs)
209 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
210 liftExpr body `thenLM` \ (body', body_info) ->
211 returnLM (StgLet (StgNonRec binder rhs') body',
212 rhs_info `unionLiftInfo` body_info)
214 | otherwise -- It's a lambda
215 = -- Do the body of the let
216 fixLM (\ ~(sc_inline, _, _) ->
217 addScInlines [binder] [sc_inline] (
219 ) `thenLM` \ (body', body_info) ->
222 dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
224 -- All occurrences in function position, so lambda lift
225 getFinalFreeVars (rhsFreeVars rhs) `thenLM` \ final_free_vars ->
227 mkScPieces final_free_vars (binder,rhs') `thenLM` \ (sc_inline, sc_bind) ->
231 nonRecScBind rhs_info sc_bind `unionLiftInfo` body_info)
233 ) `thenLM` \ (_, expr', final_info) ->
235 returnLM (expr', final_info)
237 liftExpr (StgLet (StgRec pairs) body)
239 | not (all isLiftableRec rhss)
240 = liftExpr body `thenLM` \ (body', body_info) ->
241 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
242 returnLM (StgLet (StgRec (binders `zipEqual` rhss')) body',
243 foldr unionLiftInfo body_info rhs_infos)
245 | otherwise -- All rhss are liftable
246 = -- Do the body of the let
247 fixLM (\ ~(sc_inlines, _, _) ->
248 addScInlines binders sc_inlines (
250 liftExpr body `thenLM` \ (body', body_info) ->
251 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
253 -- Find the free vars of all the rhss,
254 -- excluding the binders themselves.
255 rhs_free_vars = unionManyUniqSets (map rhsFreeVars rhss)
259 rhs_info = unionLiftInfos rhs_infos
261 getFinalFreeVars rhs_free_vars `thenLM` \ final_free_vars ->
263 mapAndUnzipLM (mkScPieces final_free_vars) (binders `zip` rhss')
264 `thenLM` \ (sc_inlines, sc_pairs) ->
265 returnLM (sc_inlines,
267 recScBind rhs_info sc_pairs `unionLiftInfo` body_info)
269 )) `thenLM` \ (_, expr', final_info) ->
271 returnLM (expr', final_info)
273 (binders,rhss) = unzip pairs
277 liftExpr (StgSCC ty cc expr)
278 = liftExpr expr `thenLM` \ (expr2, expr_info) ->
279 returnLM (StgSCC ty cc expr2, expr_info)
282 A binding is liftable if it's a *function* (args not null) and never
283 occurs in an argument position.
286 isLiftable :: PlainStgRhs -> Bool
288 isLiftable (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
289 -- experimental evidence suggests we should lift only if we will be abstracting up to 4 fvs.
290 = if not (null args || -- Not a function
291 unapplied_occ || -- Has an occ with no args at all
292 arg_occ || -- Occurs in arg position
293 length fvs > 4 -- Too many free variables
295 then {-trace ("LL: " ++ show (length fvs))-} True
297 isLiftable other_rhs = False
299 isLiftableRec :: PlainStgRhs -> Bool
300 -- this is just the same as for non-rec, except we only lift to abstract up to 1 argument
301 -- this avoids undoing Static Argument Transformation work
302 isLiftableRec (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
303 = if not (null args || -- Not a function
304 unapplied_occ || -- Has an occ with no args at all
305 arg_occ || -- Occurs in arg position
306 length fvs > 1 -- Too many free variables
308 then {-trace ("LLRec: " ++ show (length fvs))-} True
310 isLiftableRec other_rhs = False
312 rhsFreeVars :: PlainStgRhs -> IdSet
313 rhsFreeVars (StgRhsClosure _ _ fvs _ _ _) = mkUniqSet fvs
314 rhsFreeVars other = panic "rhsFreeVars"
317 dontLiftRhs is like liftExpr, except that it does not lift a top-level lambda
318 abstraction. It is used for the right-hand sides of definitions where
319 we've decided *not* to lift: for example, top-level ones or mutually-recursive
320 ones where not all are lambdas.
323 dontLiftRhs :: PlainStgRhs -> LiftM (PlainStgRhs, LiftInfo)
325 dontLiftRhs rhs@(StgRhsCon cc v args) = returnLM (rhs, emptyLiftInfo)
327 dontLiftRhs (StgRhsClosure cc bi fvs upd args body)
328 = liftExpr body `thenLM` \ (body', body_info) ->
329 returnLM (StgRhsClosure cc bi fvs upd args body', body_info)
334 mkScPieces :: IdSet -- Extra args for the supercombinator
335 -> (Id, PlainStgRhs) -- The processed RHS and original Id
336 -> LiftM ((Id,[Id]), -- Replace abstraction with this;
337 -- the set is its free vars
338 (Id,PlainStgRhs)) -- Binding for supercombinator
340 mkScPieces extra_arg_set (id, StgRhsClosure cc bi _ upd args body)
341 = ASSERT( n_args > 0 )
342 -- Construct the rhs of the supercombinator, and its Id
343 -- this trace blackholes sometimes, don't use it
344 -- trace ("LL " ++ show (length (uniqSetToList extra_arg_set))) (
345 newSupercombinator sc_ty arity `thenLM` \ sc_id ->
347 returnLM ((sc_id, extra_args), (sc_id, sc_rhs))
351 extra_args = uniqSetToList extra_arg_set
352 arity = n_args + length extra_args
354 -- Construct the supercombinator type
355 type_of_original_id = getIdUniType id
356 extra_arg_tys = map getIdUniType extra_args
357 (tyvars, rest) = splitForalls type_of_original_id
358 sc_ty = mkForallTy tyvars (glueTyArgs extra_arg_tys rest)
360 sc_rhs = StgRhsClosure cc bi [] upd (extra_args ++ args) body
364 %************************************************************************
366 \subsection[Lift-monad]{The LiftM monad}
368 %************************************************************************
370 The monad is used only to distribute global stuff, and the unique supply.
373 type LiftM a = LiftFlags
375 -> (IdEnv -- Domain = candidates for lifting
376 (Id, -- The supercombinator
377 [Id]) -- Args to apply it to
382 type LiftFlags = Maybe Int -- No of fvs reqd to float recursive
383 -- binding; Nothing == infinity
386 runLM :: LiftFlags -> SplitUniqSupply -> LiftM a -> a
387 runLM flags us m = m flags us nullIdEnv
389 thenLM :: LiftM a -> (a -> LiftM b) -> LiftM b
390 thenLM m k ci us idenv
391 = k (m ci us1 idenv) ci us2 idenv
393 (us1, us2) = splitUniqSupply us
395 returnLM :: a -> LiftM a
396 returnLM a ci us idenv = a
398 fixLM :: (a -> LiftM a) -> LiftM a
399 fixLM k ci us idenv = r
403 mapLM :: (a -> LiftM b) -> [a] -> LiftM [b]
404 mapLM f [] = returnLM []
405 mapLM f (a:as) = f a `thenLM` \ r ->
406 mapLM f as `thenLM` \ rs ->
409 mapAndUnzipLM :: (a -> LiftM (b,c)) -> [a] -> LiftM ([b],[c])
410 mapAndUnzipLM f [] = returnLM ([],[])
411 mapAndUnzipLM f (a:as) = f a `thenLM` \ (b,c) ->
412 mapAndUnzipLM f as `thenLM` \ (bs,cs) ->
413 returnLM (b:bs, c:cs)
417 newSupercombinator :: UniType
421 newSupercombinator ty arity ci us idenv
422 = (mkSysLocal SLIT("sc") uniq ty mkUnknownSrcLoc) -- ToDo: improve location
424 -- ToDo: rm the addIdArity? Just let subsequent stg-saturation pass do it?
428 lookup :: Id -> LiftM (Id,[Id])
430 = case lookupIdEnv idenv v of
431 Just result -> result
434 addScInlines :: [Id] -> [(Id,[Id])] -> LiftM a -> LiftM a
435 addScInlines ids values m ci us idenv
438 idenv' = growIdEnvList idenv (ids `zip_lazy` values)
440 -- zip_lazy zips two things together but matches lazily on the
441 -- second argument. This is important, because the ids are know here,
442 -- but the things they are bound to are decided only later
444 zip_lazy (x:xs) ~(y:ys) = (x,y) : zip_lazy xs ys
447 -- The free vars reported by the free-var analyser will include
448 -- some ids, f, which are to be replaced by ($f a b c), where $f
449 -- is the supercombinator. Hence instead of f being a free var,
454 -- f a = ...y1..y2.....
461 -- Here the free vars of g are {f,z}; but f will be lambda-lifted
462 -- with free vars {y1,y2}, so the "real~ free vars of g are {y1,y2,z}.
464 getFinalFreeVars :: IdSet -> LiftM IdSet
466 getFinalFreeVars free_vars ci us idenv
467 = unionManyUniqSets (map munge_it (uniqSetToList free_vars))
469 munge_it :: Id -> IdSet -- Takes a free var and maps it to the "real"
471 munge_it id = case lookupIdEnv idenv id of
472 Just (_, args) -> mkUniqSet args
473 Nothing -> singletonUniqSet id
478 %************************************************************************
480 \subsection[Lift-info]{The LiftInfo type}
482 %************************************************************************
485 type LiftInfo = Bag PlainStgBinding -- Float to top
487 emptyLiftInfo = emptyBag
489 unionLiftInfo :: LiftInfo -> LiftInfo -> LiftInfo
490 unionLiftInfo binds1 binds2 = binds1 `unionBags` binds2
492 unionLiftInfos :: [LiftInfo] -> LiftInfo
493 unionLiftInfos infos = foldr unionLiftInfo emptyLiftInfo infos
495 mkScInfo :: PlainStgBinding -> LiftInfo
496 mkScInfo bind = unitBag bind
498 nonRecScBind :: LiftInfo -- From body of supercombinator
499 -> (Id, PlainStgRhs) -- Supercombinator and its rhs
501 nonRecScBind binds (sc_id,sc_rhs) = binds `snocBag` (StgNonRec sc_id sc_rhs)
504 -- In the recursive case, all the SCs from the RHSs of the recursive group
505 -- are dealing with might potentially mention the new, recursive SCs.
506 -- So we flatten the whole lot into a single recursive group.
508 recScBind :: LiftInfo -- From body of supercombinator
509 -> [(Id,PlainStgRhs)] -- Supercombinator rhs
512 recScBind binds pairs = unitBag (co_rec_ify (StgRec pairs : bagToList binds))
514 co_rec_ify :: [PlainStgBinding] -> PlainStgBinding
515 co_rec_ify binds = StgRec (concat (map f binds))
517 f (StgNonRec id rhs) = [(id,rhs)]
518 f (StgRec pairs) = pairs
521 getScBinds :: LiftInfo -> [PlainStgBinding]
522 getScBinds binds = bagToList binds
524 looksLikeSATRhs [(f,StgRhsClosure _ _ _ _ ls _)] (StgApp (StgVarAtom f') args _)
525 = (f == f') && (length args == length ls)
526 looksLikeSATRhs _ _ = False