ghc/compiler/simplStg/LambdaLift.lhs

   1 %
   2 % (c) The AQUA Project, Glasgow University, 1994-1996
   3 %
   4 \section[LambdaLift]{A STG-code lambda lifter}
   5
   6 \begin{code}
   7 #include "HsVersions.h"
   8
   9 module LambdaLift ( liftProgram ) where
  10
  11 import Ubiq{-uitous-}
  12
  13 import StgSyn
  14
  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(..)
  20                         )
  21 import SrcLoc           ( mkUnknownSrcLoc )
  22 import Type             ( splitForAllTy, mkForAllTys, mkFunTys )
  23 import UniqSupply       ( getUnique, splitUniqSupply )
  24 import Util             ( zipEqual, panic, assertPanic )
  25 \end{code}
  26
  27 This is the lambda lifter.  It turns lambda abstractions into
  28 supercombinators on a selective basis:
  29
  30 * Let-no-escaped bindings are never lifted. That's one major reason
  31   why the lambda lifter is done in STG.
  32
  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:
  36
  37         let
  38           f = \x -> e
  39         in
  40         ..(f a1)...(f a2)...
  41   thus
  42
  43     $f p q r x = e      -- Supercombinator
  44
  45         ..($f p q r a1)...($f p q r a2)...
  46
  47   NOTE that the original binding is eliminated.
  48
  49   But in this case, f won't be lifted:
  50
  51         let
  52           f = \x -> e
  53         in
  54         ..(g f)...(f a2)...
  55
  56   Why? Because we have to heap-allocate a closure for f thus:
  57
  58     $f p q r x = e      -- Supercombinator
  59
  60         let
  61           f = $f p q r
  62         in
  63         ..(g f)...($f p q r a2)..
  64
  65   so it might as well be the original lambda abstraction.
  66
  67   We also do not lift if the function has an occurrence with no arguments, e.g.
  68
  69         let
  70           f = \x -> e
  71         in f
  72
  73   as this form is more efficient than if we create a partial application
  74
  75   $f p q r x = e      -- Supercombinator
  76
  77         f p q r
  78
  79 * Recursive bindings *all* of whose RHSs are lambda abstractions are
  80   lifted iff
  81         - all the occurrences of all the binders are in a function position
  82         - there aren't ``too many'' free variables.
  83
  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.
  87
  88 Recent Observations:
  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...
  98                     in f
  99   ==>
 100   f = \y x -> ...y...
 101   h = \ x -> ... let y = ...
 102                  in f y
 103
 104   now f y is a partial application, so it will be updated, and this
 105   is Bad.
 106
 107
 108 --- NOT RELEVANT FOR STG ----
 109 * All ``lone'' lambda abstractions are lifted.  Notably this means lambda
 110   abstractions:
 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 -----------------------------
 114
 115 %************************************************************************
 116 %*                                                                      *
 117 \subsection[Lift-expressions]{The main function: liftExpr}
 118 %*                                                                      *
 119 %************************************************************************
 120
 121 \begin{code}
 122 liftProgram :: UniqSupply -> [StgBinding] -> [StgBinding]
 123 liftProgram us prog = concat (runLM Nothing us (mapLM liftTopBind prog))
 124
 125
 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'])
 130
 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))
 135              ])
 136   where
 137    (ids, rhss) = unzip pairs
 138 \end{code}
 139
 140
 141 \begin{code}
 142 liftExpr :: StgExpr
 143          -> LiftM (StgExpr, LiftInfo)
 144
 145
 146 liftExpr expr@(StgCon con args lvs) = returnLM (expr, emptyLiftInfo)
 147 liftExpr expr@(StgPrim op args lvs) = returnLM (expr, emptyLiftInfo)
 148
 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,
 154               emptyLiftInfo)
 155         -- The lvs field is probably wrong, but we reconstruct it
 156         -- anyway following lambda lifting
 157
 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)
 162   where
 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)
 167
 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)
 172
 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)
 176
 177     lift_prim_alt (lit, rhs)
 178         = liftExpr rhs  `thenLM` \ (rhs', rhs_info) ->
 179           returnLM ((lit, rhs'), rhs_info)
 180
 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)
 185 \end{code}
 186
 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
 189 lambda anyway.
 190
 191 \begin{code}
 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)
 197
 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)
 203   where
 204    (binders,rhss) = unzip pairs
 205 \end{code}
 206
 207 \begin{code}
 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)
 214
 215   | otherwise   -- It's a lambda
 216   =     -- Do the body of the let
 217     fixLM (\ ~(sc_inline, _, _) ->
 218       addScInlines [binder] [sc_inline] (
 219         liftExpr body
 220       )                 `thenLM` \ (body', body_info) ->
 221
 222         -- Deal with the RHS
 223       dontLiftRhs rhs           `thenLM` \ (rhs', rhs_info) ->
 224
 225         -- All occurrences in function position, so lambda lift
 226       getFinalFreeVars (rhsFreeVars rhs)    `thenLM` \ final_free_vars ->
 227
 228       mkScPieces final_free_vars (binder,rhs')  `thenLM` \ (sc_inline, sc_bind) ->
 229
 230       returnLM (sc_inline,
 231                 body',
 232                 nonRecScBind rhs_info sc_bind `unionLiftInfo` body_info)
 233
 234     )                   `thenLM` \ (_, expr', final_info) ->
 235
 236     returnLM (expr', final_info)
 237
 238 liftExpr (StgLet (StgRec pairs) body)
 239 --[Andre-testing]
 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)
 245
 246   | otherwise   -- All rhss are liftable
 247   = -- Do the body of the let
 248     fixLM (\ ~(sc_inlines, _, _) ->
 249       addScInlines binders sc_inlines   (
 250
 251       liftExpr body                     `thenLM` \ (body', body_info) ->
 252       mapAndUnzipLM dontLiftRhs rhss    `thenLM` \ (rhss', rhs_infos) ->
 253       let
 254         -- Find the free vars of all the rhss,
 255         -- excluding the binders themselves.
 256         rhs_free_vars = unionManyIdSets (map rhsFreeVars rhss)
 257                         `minusIdSet`
 258                         mkIdSet binders
 259
 260         rhs_info      = unionLiftInfos rhs_infos
 261       in
 262       getFinalFreeVars rhs_free_vars    `thenLM` \ final_free_vars ->
 263
 264       mapAndUnzipLM (mkScPieces final_free_vars) (binders `zip` rhss')
 265                                         `thenLM` \ (sc_inlines, sc_pairs) ->
 266       returnLM (sc_inlines,
 267                 body',
 268                 recScBind rhs_info sc_pairs `unionLiftInfo` body_info)
 269
 270     ))                  `thenLM` \ (_, expr', final_info) ->
 271
 272     returnLM (expr', final_info)
 273   where
 274     (binders,rhss)    = unzip pairs
 275 \end{code}
 276
 277 \begin{code}
 278 liftExpr (StgSCC ty cc expr)
 279   = liftExpr expr `thenLM` \ (expr2, expr_info) ->
 280     returnLM (StgSCC ty cc expr2, expr_info)
 281 \end{code}
 282
 283 A binding is liftable if it's a *function* (args not null) and never
 284 occurs in an argument position.
 285
 286 \begin{code}
 287 isLiftable :: StgRhs -> Bool
 288
 289 isLiftable (StgRhsClosure _ (StgBinderInfo arg_occ _ _ _ unapplied_occ) fvs _ args _)
 290
 291   -- Experimental evidence suggests we should lift only if we will be
 292   -- abstracting up to 4 fvs.
 293
 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
 298         )
 299     then {-trace ("LL: " ++ show (length fvs))-} True
 300     else False
 301 isLiftable other_rhs = False
 302
 303 isLiftableRec :: StgRhs -> Bool
 304
 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
 308
 309 {- Andre's longer comment about isLiftableRec: 1996/01:
 310
 311 A rec binding is "liftable" (according to our heuristics) if:
 312 * It is a function,
 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
 316   itself!)
 317
 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
 326 here).
 327 -}
 328
 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
 334         )
 335     then {-trace ("LLRec: " ++ show (length fvs))-} True
 336     else False
 337 isLiftableRec other_rhs = False
 338
 339 rhsFreeVars :: StgRhs -> IdSet
 340 rhsFreeVars (StgRhsClosure _ _ fvs _ _ _) = mkIdSet fvs
 341 rhsFreeVars other                         = panic "rhsFreeVars"
 342 \end{code}
 343
 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.
 348
 349 \begin{code}
 350 dontLiftRhs :: StgRhs -> LiftM (StgRhs, LiftInfo)
 351
 352 dontLiftRhs rhs@(StgRhsCon cc v args) = returnLM (rhs, emptyLiftInfo)
 353
 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)
 357 \end{code}
 358
 359 \begin{code}
 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
 365
 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))
 371   where
 372     n_args     = length args
 373     extra_args = idSetToList extra_arg_set
 374     arity      = n_args + length extra_args
 375
 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)
 381
 382     sc_rhs = StgRhsClosure cc bi [] upd (extra_args ++ args) body
 383 \end{code}
 384
 385
 386 %************************************************************************
 387 %*                                                                      *
 388 \subsection[Lift-monad]{The LiftM monad}
 389 %*                                                                      *
 390 %************************************************************************
 391
 392 The monad is used only to distribute global stuff, and the unique supply.
 393
 394 \begin{code}
 395 type LiftM a =  LiftFlags
 396              -> UniqSupply
 397              -> (IdEnv                          -- Domain = candidates for lifting
 398                        (Id,                     -- The supercombinator
 399                         [Id])                   -- Args to apply it to
 400                  )
 401              -> a
 402
 403
 404 type LiftFlags = Maybe Int      -- No of fvs reqd to float recursive
 405                                 -- binding; Nothing == infinity
 406
 407
 408 runLM :: LiftFlags -> UniqSupply -> LiftM a -> a
 409 runLM flags us m = m flags us nullIdEnv
 410
 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
 414   where
 415     (us1, us2) = splitUniqSupply us
 416
 417 returnLM :: a -> LiftM a
 418 returnLM a ci us idenv = a
 419
 420 fixLM :: (a -> LiftM a) -> LiftM a
 421 fixLM k ci us idenv = r
 422                        where
 423                          r = k r ci us idenv
 424
 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 ->
 429                  returnLM (r:rs)
 430
 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)
 436 \end{code}
 437
 438 \begin{code}
 439 newSupercombinator :: Type
 440                    -> Int               -- Arity
 441                    -> LiftM Id
 442
 443 newSupercombinator ty arity ci us idenv
 444   = (mkSysLocal SLIT("sc") uniq ty mkUnknownSrcLoc)     -- ToDo: improve location
 445     `addIdArity` arity
 446         -- ToDo: rm the addIdArity?  Just let subsequent stg-saturation pass do it?
 447   where
 448     uniq = getUnique us
 449
 450 lookup :: Id -> LiftM (Id,[Id])
 451 lookup v ci us idenv
 452   = case (lookupIdEnv idenv v) of
 453       Just result -> result
 454       Nothing     -> (v, [])
 455
 456 addScInlines :: [Id] -> [(Id,[Id])] -> LiftM a -> LiftM a
 457 addScInlines ids values m ci us idenv
 458   = m ci us idenv'
 459   where
 460     idenv' = growIdEnvList idenv (ids `zip_lazy` values)
 461
 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
 465     zip_lazy [] _           = []
 466     zip_lazy (x:xs) ~(y:ys) = (x,y) : zip_lazy xs ys
 467
 468
 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,
 472 -- {a,b,c} are.
 473 --
 474 -- Example
 475 --      let
 476 --         f a = ...y1..y2.....
 477 --      in
 478 --      let
 479 --         g b = ...f...z...
 480 --      in
 481 --      ...
 482 --
 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}.
 485
 486 getFinalFreeVars :: IdSet -> LiftM IdSet
 487
 488 getFinalFreeVars free_vars ci us idenv
 489   = unionManyIdSets (map munge_it (idSetToList free_vars))
 490   where
 491     munge_it :: Id -> IdSet     -- Takes a free var and maps it to the "real"
 492                                 -- free var
 493     munge_it id = case (lookupIdEnv idenv id) of
 494                     Just (_, args) -> mkIdSet args
 495                     Nothing        -> unitIdSet id
 496 \end{code}
 497
 498
 499 %************************************************************************
 500 %*                                                                      *
 501 \subsection[Lift-info]{The LiftInfo type}
 502 %*                                                                      *
 503 %************************************************************************
 504
 505 \begin{code}
 506 type LiftInfo = Bag StgBinding  -- Float to top
 507
 508 emptyLiftInfo = emptyBag
 509
 510 unionLiftInfo :: LiftInfo -> LiftInfo -> LiftInfo
 511 unionLiftInfo binds1 binds2 = binds1 `unionBags` binds2
 512
 513 unionLiftInfos :: [LiftInfo] -> LiftInfo
 514 unionLiftInfos infos = foldr unionLiftInfo emptyLiftInfo infos
 515
 516 mkScInfo :: StgBinding -> LiftInfo
 517 mkScInfo bind = unitBag bind
 518
 519 nonRecScBind :: LiftInfo                -- From body of supercombinator
 520              -> (Id, StgRhs)    -- Supercombinator and its rhs
 521              -> LiftInfo
 522 nonRecScBind binds (sc_id,sc_rhs) = binds `snocBag` (StgNonRec sc_id sc_rhs)
 523
 524
 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.
 528
 529 recScBind :: LiftInfo                   -- From body of supercombinator
 530            -> [(Id,StgRhs)]     -- Supercombinator rhs
 531            -> LiftInfo
 532
 533 recScBind binds pairs = unitBag (co_rec_ify (StgRec pairs : bagToList binds))
 534
 535 co_rec_ify :: [StgBinding] -> StgBinding
 536 co_rec_ify binds = StgRec (concat (map f binds))
 537   where
 538     f (StgNonRec id rhs) = [(id,rhs)]
 539     f (StgRec pairs)     = pairs
 540
 541
 542 getScBinds :: LiftInfo -> [StgBinding]
 543 getScBinds binds = bagToList binds
 544
 545 looksLikeSATRhs [(f,StgRhsClosure _ _ _ _ ls _)] (StgApp (StgVarArg f') args _)
 546   = (f == f') && (length args == length ls)
 547 looksLikeSATRhs _ _ = False
 548 \end{code}