2 % (c) The AQUA Project, Glasgow University, 1994-1998
4 \section[LambdaLift]{A STG-code lambda lifter}
7 module LambdaLift ( liftProgram ) where
9 #include "HsVersions.h"
13 import CmdLineOpts ( opt_EnsureSplittableC )
14 import Bag ( Bag, emptyBag, unionBags, unitBag, snocBag, bagToList )
15 import Id ( mkVanillaId, idType, setIdArityInfo, Id )
18 import IdInfo ( exactArity )
19 import Module ( Module )
20 import Name ( Name, mkGlobalName, mkLocalName )
21 import OccName ( mkVarOcc )
22 import Type ( splitForAllTys, mkForAllTys, mkFunTys, Type )
23 import Unique ( Unique )
24 import UniqSupply ( uniqFromSupply, splitUniqSupply, UniqSupply )
25 import Util ( zipEqual )
26 import SrcLoc ( noSrcLoc )
27 import Panic ( panic, assertPanic )
30 This is the lambda lifter. It turns lambda abstractions into
31 supercombinators on a selective basis:
33 * Let-no-escaped bindings are never lifted. That's one major reason
34 why the lambda lifter is done in STG.
36 * Non-recursive bindings whose RHS is a lambda abstractions are lifted,
37 provided all the occurrences of the bound variable is in a function
38 postition. In this example, f will be lifted:
46 $f p q r x = e -- Supercombinator
48 ..($f p q r a1)...($f p q r a2)...
50 NOTE that the original binding is eliminated.
52 But in this case, f won't be lifted:
59 Why? Because we have to heap-allocate a closure for f thus:
61 $f p q r x = e -- Supercombinator
66 ..(g f)...($f p q r a2)..
68 so it might as well be the original lambda abstraction.
70 We also do not lift if the function has an occurrence with no arguments, e.g.
76 as this form is more efficient than if we create a partial application
78 $f p q r x = e -- Supercombinator
82 * Recursive bindings *all* of whose RHSs are lambda abstractions are
84 - all the occurrences of all the binders are in a function position
85 - there aren't ``too many'' free variables.
87 Same reasoning as before for the function-position stuff. The ``too many
88 free variable'' part comes from considering the (potentially many)
89 recursive calls, which may now have lots of free vars.
93 * 2 might be already ``too many'' variables to abstract.
94 The problem is that the increase in the number of free variables
95 of closures refering to the lifted function (which is always # of
96 abstracted args - 1) may increase heap allocation a lot.
97 Expeiments are being done to check this...
99 * We do not lambda lift if the function has at least one occurrence
100 without any arguments. This caused lots of problems. Ex:
101 h = \ x -> ... let y = ...
102 in let let f = \x -> ...y...
106 h = \ x -> ... let y = ...
109 now f y is a partial application, so it will be updated, and this
113 --- NOT RELEVANT FOR STG ----
114 * All ``lone'' lambda abstractions are lifted. Notably this means lambda
116 - in a case alternative: case e of True -> (\x->b)
117 - in the body of a let: let x=e in (\y->b)
118 -----------------------------
120 %************************************************************************
122 \subsection[Lift-expressions]{The main function: liftExpr}
124 %************************************************************************
127 liftProgram :: Module -> UniqSupply -> [StgBinding] -> [StgBinding]
128 liftProgram mod us prog = concat (runLM mod Nothing us (mapLM liftTopBind prog))
131 liftTopBind :: StgBinding -> LiftM [StgBinding]
132 liftTopBind (StgNonRec id rhs)
133 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
134 returnLM (getScBinds rhs_info ++ [StgNonRec id rhs'])
136 liftTopBind (StgRec pairs)
137 = mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
138 returnLM ([co_rec_ify (StgRec (ids `zip` rhss') :
139 getScBinds (unionLiftInfos rhs_infos))
142 (ids, rhss) = unzip pairs
148 -> LiftM (StgExpr, LiftInfo)
151 liftExpr expr@(StgLit _) = returnLM (expr, emptyLiftInfo)
152 liftExpr expr@(StgConApp _ _) = returnLM (expr, emptyLiftInfo)
153 liftExpr expr@(StgPrimApp _ _ _) = returnLM (expr, emptyLiftInfo)
155 liftExpr expr@(StgApp v args)
156 = lookUp v `thenLM` \ ~(sc, sc_args) -> -- NB the ~. We don't want to
157 -- poke these bindings too early!
158 returnLM (StgApp sc (map StgVarArg sc_args ++ args),
160 -- The lvs field is probably wrong, but we reconstruct it
161 -- anyway following lambda lifting
163 liftExpr (StgCase scrut lv1 lv2 bndr srt alts)
164 = liftExpr scrut `thenLM` \ (scrut', scrut_info) ->
165 lift_alts alts `thenLM` \ (alts', alts_info) ->
166 returnLM (StgCase scrut' lv1 lv2 bndr srt alts', scrut_info `unionLiftInfo` alts_info)
168 lift_alts (StgAlgAlts tycon alg_alts deflt)
169 = mapAndUnzipLM lift_alg_alt alg_alts `thenLM` \ (alg_alts', alt_infos) ->
170 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
171 returnLM (StgAlgAlts tycon alg_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
173 lift_alts (StgPrimAlts tycon prim_alts deflt)
174 = mapAndUnzipLM lift_prim_alt prim_alts `thenLM` \ (prim_alts', alt_infos) ->
175 lift_deflt deflt `thenLM` \ (deflt', deflt_info) ->
176 returnLM (StgPrimAlts tycon prim_alts' deflt', foldr unionLiftInfo deflt_info alt_infos)
178 lift_alg_alt (con, args, use_mask, rhs)
179 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
180 returnLM ((con, args, use_mask, rhs'), rhs_info)
182 lift_prim_alt (lit, rhs)
183 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
184 returnLM ((lit, rhs'), rhs_info)
186 lift_deflt StgNoDefault = returnLM (StgNoDefault, emptyLiftInfo)
187 lift_deflt (StgBindDefault rhs)
188 = liftExpr rhs `thenLM` \ (rhs', rhs_info) ->
189 returnLM (StgBindDefault rhs', rhs_info)
192 Now the interesting cases. Let no escape isn't lifted. We turn it
193 back into a let, to play safe, because we have to redo that pass after
197 liftExpr (StgLetNoEscape _ _ (StgNonRec binder rhs) body)
198 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
199 liftExpr body `thenLM` \ (body', body_info) ->
200 returnLM (StgLet (StgNonRec binder rhs') body',
201 rhs_info `unionLiftInfo` body_info)
203 liftExpr (StgLetNoEscape _ _ (StgRec pairs) body)
204 = liftExpr body `thenLM` \ (body', body_info) ->
205 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
206 returnLM (StgLet (StgRec (zipEqual "liftExpr" binders rhss')) body',
207 foldr unionLiftInfo body_info rhs_infos)
209 (binders,rhss) = unzip pairs
213 liftExpr (StgLet (StgNonRec binder rhs) body)
214 | not (isLiftable rhs)
215 = dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
216 liftExpr body `thenLM` \ (body', body_info) ->
217 returnLM (StgLet (StgNonRec binder rhs') body',
218 rhs_info `unionLiftInfo` body_info)
220 | otherwise -- It's a lambda
221 = -- Do the body of the let
222 fixLM (\ ~(sc_inline, _, _) ->
223 addScInlines [binder] [sc_inline] (
225 ) `thenLM` \ (body', body_info) ->
228 dontLiftRhs rhs `thenLM` \ (rhs', rhs_info) ->
230 -- All occurrences in function position, so lambda lift
231 getFinalFreeVars (rhsFreeVars rhs) `thenLM` \ final_free_vars ->
233 mkScPieces final_free_vars (binder,rhs') `thenLM` \ (sc_inline, sc_bind) ->
237 nonRecScBind rhs_info sc_bind `unionLiftInfo` body_info)
239 ) `thenLM` \ (_, expr', final_info) ->
241 returnLM (expr', final_info)
243 liftExpr (StgLet (StgRec pairs) body)
245 | not (all isLiftableRec rhss)
246 = liftExpr body `thenLM` \ (body', body_info) ->
247 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
248 returnLM (StgLet (StgRec (zipEqual "liftExpr2" binders rhss')) body',
249 foldr unionLiftInfo body_info rhs_infos)
251 | otherwise -- All rhss are liftable
252 = -- Do the body of the let
253 fixLM (\ ~(sc_inlines, _, _) ->
254 addScInlines binders sc_inlines (
256 liftExpr body `thenLM` \ (body', body_info) ->
257 mapAndUnzipLM dontLiftRhs rhss `thenLM` \ (rhss', rhs_infos) ->
259 -- Find the free vars of all the rhss,
260 -- excluding the binders themselves.
261 rhs_free_vars = unionVarSets (map rhsFreeVars rhss)
265 rhs_info = unionLiftInfos rhs_infos
267 getFinalFreeVars rhs_free_vars `thenLM` \ final_free_vars ->
269 mapAndUnzipLM (mkScPieces final_free_vars) (binders `zip` rhss')
270 `thenLM` \ (sc_inlines, sc_pairs) ->
271 returnLM (sc_inlines,
273 recScBind rhs_info sc_pairs `unionLiftInfo` body_info)
275 )) `thenLM` \ (_, expr', final_info) ->
277 returnLM (expr', final_info)
279 (binders,rhss) = unzip pairs
283 liftExpr (StgSCC cc expr)
284 = liftExpr expr `thenLM` \ (expr2, expr_info) ->
285 returnLM (StgSCC cc expr2, expr_info)
288 A binding is liftable if it's a *function* (args not null) and never
289 occurs in an argument position.
292 isLiftable :: StgRhs -> Bool
294 isLiftable (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) _ fvs _ args _)
296 -- Experimental evidence suggests we should lift only if we will be
297 -- abstracting up to 4 fvs.
299 = if not (null args || -- Not a function
300 unapplied_occ || -- Has an occ with no args at all
301 arg_occ || -- Occurs in arg position
302 length fvs > 4 -- Too many free variables
304 then {-trace ("LL: " ++ show (length fvs))-} True
306 isLiftable other_rhs = False
308 isLiftableRec :: StgRhs -> Bool
310 -- this is just the same as for non-rec, except we only lift to
311 -- abstract up to 1 argument this avoids undoing Static Argument
312 -- Transformation work
314 {- Andre's longer comment about isLiftableRec: 1996/01:
316 A rec binding is "liftable" (according to our heuristics) if:
318 * all occurrences have arguments,
319 * does not occur in an argument position and
320 * has up to *2* free variables (including the rec binding variable
323 The point is: my experiments show that SAT is more important than LL.
324 Therefore if we still want to do LL, for *recursive* functions, we do
325 not want LL to undo what SAT did. We do this by avoiding LL recursive
326 functions that have more than 2 fvs, since if this recursive function
327 was created by SAT (we don't know!), it would have at least 3 fvs: one
328 for the rec binding itself and 2 more for the static arguments (note:
329 this matches with the choice of performing SAT to have at least 2
330 static arguments, if we change things there we should change things
334 isLiftableRec (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) _ fvs _ args _)
335 = if not (null args || -- Not a function
336 unapplied_occ || -- Has an occ with no args at all
337 arg_occ || -- Occurs in arg position
338 length fvs > 2 -- Too many free variables
340 then {-trace ("LLRec: " ++ show (length fvs))-} True
342 isLiftableRec other_rhs = False
344 rhsFreeVars :: StgRhs -> IdSet
345 rhsFreeVars (StgRhsClosure _ _ _ fvs _ _ _) = mkVarSet fvs
346 rhsFreeVars other = panic "rhsFreeVars"
349 dontLiftRhs is like liftExpr, except that it does not lift a top-level
350 lambda abstraction. It is used for the right-hand sides of
351 definitions where we've decided *not* to lift: for example, top-level
352 ones or mutually-recursive ones where not all are lambdas.
355 dontLiftRhs :: StgRhs -> LiftM (StgRhs, LiftInfo)
357 dontLiftRhs rhs@(StgRhsCon cc v args) = returnLM (rhs, emptyLiftInfo)
359 dontLiftRhs (StgRhsClosure cc bi srt fvs upd args body)
360 = liftExpr body `thenLM` \ (body', body_info) ->
361 returnLM (StgRhsClosure cc bi srt fvs upd args body', body_info)
365 mkScPieces :: IdSet -- Extra args for the supercombinator
366 -> (Id, StgRhs) -- The processed RHS and original Id
367 -> LiftM ((Id,[Id]), -- Replace abstraction with this;
368 -- the set is its free vars
369 (Id,StgRhs)) -- Binding for supercombinator
371 mkScPieces extra_arg_set (id, StgRhsClosure cc bi srt _ upd args body)
372 = ASSERT( n_args > 0 )
373 -- Construct the rhs of the supercombinator, and its Id
374 newSupercombinator sc_ty arity `thenLM` \ sc_id ->
375 returnLM ((sc_id, extra_args), (sc_id, sc_rhs))
378 extra_args = varSetElems extra_arg_set
379 arity = n_args + length extra_args
381 -- Construct the supercombinator type
382 type_of_original_id = idType id
383 extra_arg_tys = map idType extra_args
384 (tyvars, rest) = splitForAllTys type_of_original_id
385 sc_ty = mkForAllTys tyvars (mkFunTys extra_arg_tys rest)
387 sc_rhs = StgRhsClosure cc bi srt [] upd (extra_args ++ args) body
391 %************************************************************************
393 \subsection[Lift-monad]{The LiftM monad}
395 %************************************************************************
397 The monad is used only to distribute global stuff, and the unique supply.
400 type LiftM a = Module
403 -> (IdEnv -- Domain = candidates for lifting
404 (Id, -- The supercombinator
405 [Id]) -- Args to apply it to
410 type LiftFlags = Maybe Int -- No of fvs reqd to float recursive
411 -- binding; Nothing == infinity
414 runLM :: Module -> LiftFlags -> UniqSupply -> LiftM a -> a
415 runLM mod flags us m = m mod flags us emptyVarEnv
417 thenLM :: LiftM a -> (a -> LiftM b) -> LiftM b
418 thenLM m k mod ci us idenv
419 = k (m mod ci us1 idenv) mod ci us2 idenv
421 (us1, us2) = splitUniqSupply us
423 returnLM :: a -> LiftM a
424 returnLM a mod ci us idenv = a
426 fixLM :: (a -> LiftM a) -> LiftM a
427 fixLM k mod ci us idenv = r
429 r = k r mod ci us idenv
431 mapLM :: (a -> LiftM b) -> [a] -> LiftM [b]
432 mapLM f [] = returnLM []
433 mapLM f (a:as) = f a `thenLM` \ r ->
434 mapLM f as `thenLM` \ rs ->
437 mapAndUnzipLM :: (a -> LiftM (b,c)) -> [a] -> LiftM ([b],[c])
438 mapAndUnzipLM f [] = returnLM ([],[])
439 mapAndUnzipLM f (a:as) = f a `thenLM` \ (b,c) ->
440 mapAndUnzipLM f as `thenLM` \ (bs,cs) ->
441 returnLM (b:bs, c:cs)
445 newSupercombinator :: Type
449 newSupercombinator ty arity mod ci us idenv
450 = mkVanillaId (mkTopName uniq mod SLIT("_ll")) ty
451 `setIdArityInfo` exactArity arity
452 -- ToDo: rm the setIdArity? Just let subsequent stg-saturation pass do it?
454 uniq = uniqFromSupply us
457 mkTopName :: Unique -> Module -> FAST_STRING -> Name
458 -- Make a top-level name; make it Global if top-level
459 -- things should be externally visible; Local otherwise
460 -- This chap is only used *after* the tidyCore phase
461 -- Notably, it is used during STG lambda lifting
463 -- We have to make sure that the name is globally unique
464 -- and we don't have tidyCore to help us. So we append
465 -- the unique. Hack! Hack!
466 -- (Used only by the STG lambda lifter.)
467 mkTopName uniq mod fs
468 | opt_EnsureSplittableC = mkGlobalName uniq mod occ noSrcLoc
469 | otherwise = mkLocalName uniq occ noSrcLoc
471 occ = mkVarOcc (_PK_ ((_UNPK_ fs) ++ show uniq))
473 lookUp :: Id -> LiftM (Id,[Id])
474 lookUp v mod ci us idenv
475 = case (lookupVarEnv idenv v) of
476 Just result -> result
479 addScInlines :: [Id] -> [(Id,[Id])] -> LiftM a -> LiftM a
480 addScInlines ids values m mod ci us idenv
483 idenv' = extendVarEnvList idenv (ids `zip_lazy` values)
485 -- zip_lazy zips two things together but matches lazily on the
486 -- second argument. This is important, because the ids are know here,
487 -- but the things they are bound to are decided only later
489 zip_lazy (x:xs) ~(y:ys) = (x,y) : zip_lazy xs ys
492 -- The free vars reported by the free-var analyser will include
493 -- some ids, f, which are to be replaced by ($f a b c), where $f
494 -- is the supercombinator. Hence instead of f being a free var,
499 -- f a = ...y1..y2.....
506 -- Here the free vars of g are {f,z}; but f will be lambda-lifted
507 -- with free vars {y1,y2}, so the "real~ free vars of g are {y1,y2,z}.
509 getFinalFreeVars :: IdSet -> LiftM IdSet
511 getFinalFreeVars free_vars mod ci us idenv
512 = unionVarSets (map munge_it (varSetElems free_vars))
514 munge_it :: Id -> IdSet -- Takes a free var and maps it to the "real"
516 munge_it id = case (lookupVarEnv idenv id) of
517 Just (_, args) -> mkVarSet args
518 Nothing -> unitVarSet id
522 %************************************************************************
524 \subsection[Lift-info]{The LiftInfo type}
526 %************************************************************************
529 type LiftInfo = Bag StgBinding -- Float to top
531 emptyLiftInfo = emptyBag
533 unionLiftInfo :: LiftInfo -> LiftInfo -> LiftInfo
534 unionLiftInfo binds1 binds2 = binds1 `unionBags` binds2
536 unionLiftInfos :: [LiftInfo] -> LiftInfo
537 unionLiftInfos infos = foldr unionLiftInfo emptyLiftInfo infos
539 mkScInfo :: StgBinding -> LiftInfo
540 mkScInfo bind = unitBag bind
542 nonRecScBind :: LiftInfo -- From body of supercombinator
543 -> (Id, StgRhs) -- Supercombinator and its rhs
545 nonRecScBind binds (sc_id,sc_rhs) = binds `snocBag` (StgNonRec sc_id sc_rhs)
548 -- In the recursive case, all the SCs from the RHSs of the recursive group
549 -- are dealing with might potentially mention the new, recursive SCs.
550 -- So we flatten the whole lot into a single recursive group.
552 recScBind :: LiftInfo -- From body of supercombinator
553 -> [(Id,StgRhs)] -- Supercombinator rhs
556 recScBind binds pairs = unitBag (co_rec_ify (StgRec pairs : bagToList binds))
558 co_rec_ify :: [StgBinding] -> StgBinding
559 co_rec_ify binds = StgRec (concat (map f binds))
561 f (StgNonRec id rhs) = [(id,rhs)]
562 f (StgRec pairs) = pairs
565 getScBinds :: LiftInfo -> [StgBinding]
566 getScBinds binds = bagToList binds