ghc/compiler/simplStg/StgVarInfo.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
   3 %
   4 \section[StgVarInfo]{Sets free/live variable info in STG syntax}
   5
   6 And, as we have the info in hand, we may convert some lets to
   7 let-no-escapes.
   8
   9 \begin{code}
  10 module StgVarInfo ( setStgVarInfo ) where
  11
  12 #include "HsVersions.h"
  13
  14 import StgSyn
  15
  16 import Id               ( setIdArity, getIdArity, Id )
  17 import VarSet
  18 import VarEnv
  19 import Var
  20 import IdInfo           ( ArityInfo(..), InlinePragInfo(..),
  21                           setInlinePragInfo )
  22 import Maybes           ( maybeToBool )
  23 import Name             ( isLocallyDefined )
  24 import BasicTypes       ( Arity )
  25 import Outputable
  26
  27 infixr 9 `thenLne`, `thenLne_`
  28 \end{code}
  29
  30 %************************************************************************
  31 %*                                                                      *
  32 \subsection[live-vs-free-doc]{Documentation}
  33 %*                                                                      *
  34 %************************************************************************
  35
  36 (There is other relevant documentation in codeGen/CgLetNoEscape.)
  37
  38 March 97: setStgVarInfo guarantees to leave every variable's arity correctly
  39 set.  The lambda lifter makes some let-bound variables (which have arities)
  40 and turns them into lambda-bound ones (which should not, else we get Vap trouble),
  41 so this guarantee is necessary, as well as desirable.
  42
  43 The arity information is used in the code generator, when deciding if
  44 a right-hand side is a saturated application so we can generate a VAP
  45 closure.
  46
  47 The actual Stg datatype is decorated with {\em live variable}
  48 information, as well as {\em free variable} information.  The two are
  49 {\em not} the same.  Liveness is an operational property rather than a
  50 semantic one.  A variable is live at a particular execution point if
  51 it can be referred to {\em directly} again.  In particular, a dead
  52 variable's stack slot (if it has one):
  53 \begin{enumerate}
  54 \item
  55 should be stubbed to avoid space leaks, and
  56 \item
  57 may be reused for something else.
  58 \end{enumerate}
  59
  60 There ought to be a better way to say this.  Here are some examples:
  61 \begin{verbatim}
  62         let v = [q] \[x] -> e
  63         in
  64         ...v...  (but no q's)
  65 \end{verbatim}
  66
  67 Just after the `in', v is live, but q is dead.  If the whole of that
  68 let expression was enclosed in a case expression, thus:
  69 \begin{verbatim}
  70         case (let v = [q] \[x] -> e in ...v...) of
  71                 alts[...q...]
  72 \end{verbatim}
  73 (ie @alts@ mention @q@), then @q@ is live even after the `in'; because
  74 we'll return later to the @alts@ and need it.
  75
  76 Let-no-escapes make this a bit more interesting:
  77 \begin{verbatim}
  78         let-no-escape v = [q] \ [x] -> e
  79         in
  80         ...v...
  81 \end{verbatim}
  82 Here, @q@ is still live at the `in', because @v@ is represented not by
  83 a closure but by the current stack state.  In other words, if @v@ is
  84 live then so is @q@.  Furthermore, if @e@ mentions an enclosing
  85 let-no-escaped variable, then {\em its} free variables are also live
  86 if @v@ is.
  87
  88 %************************************************************************
  89 %*                                                                      *
  90 \subsection[binds-StgVarInfo]{Setting variable info: top-level, binds, RHSs}
  91 %*                                                                      *
  92 %************************************************************************
  93
  94 Top-level:
  95 \begin{code}
  96 setStgVarInfo   :: Bool                 -- True <=> do let-no-escapes
  97                 -> [StgBinding] -- input
  98                 -> [StgBinding] -- result
  99
 100 setStgVarInfo want_LNEs pgm
 101   = pgm'
 102   where
 103     (pgm', _) = initLne want_LNEs (varsTopBinds pgm)
 104
 105 \end{code}
 106
 107 For top-level guys, we basically aren't worried about this
 108 live-variable stuff; we do need to keep adding to the environment
 109 as we step through the bindings (using @extendVarEnv@).
 110
 111 \begin{code}
 112 varsTopBinds :: [StgBinding] -> LneM ([StgBinding], FreeVarsInfo)
 113
 114 varsTopBinds [] = returnLne ([], emptyFVInfo)
 115 varsTopBinds (bind:binds)
 116   = extendVarEnvLne env_extension (
 117         varsTopBinds binds                      `thenLne` \ (binds', fv_binds) ->
 118         varsTopBind binders' fv_binds bind      `thenLne` \ (bind',  fv_bind) ->
 119         returnLne ((bind' : binds'),
 120                    (fv_binds `unionFVInfo` fv_bind) `minusFVBinders` binders'
 121                   )
 122
 123     )
 124   where
 125     pairs         = case bind of
 126                         StgNonRec binder rhs -> [(binder,rhs)]
 127                         StgRec pairs         -> pairs
 128
 129     binders' = [ binder `setIdArity` ArityExactly (rhsArity rhs)
 130                | (binder, rhs) <- pairs
 131                ]
 132
 133     env_extension = binders' `zip` repeat how_bound
 134
 135     how_bound = LetrecBound
 136                         True {- top level -}
 137                         emptyVarSet
 138
 139
 140 varsTopBind :: [Id]                     -- New binders (with correct arity)
 141             -> FreeVarsInfo             -- Info about the body
 142             -> StgBinding
 143             -> LneM (StgBinding, FreeVarsInfo)
 144
 145 varsTopBind [binder'] body_fvs (StgNonRec binder rhs)
 146   = varsRhs body_fvs (binder,rhs)               `thenLne` \ (rhs2, fvs, _) ->
 147     returnLne (StgNonRec binder' rhs2, fvs)
 148
 149 varsTopBind binders' body_fvs (StgRec pairs)
 150   = fixLne (\ ~(_, rec_rhs_fvs) ->
 151         let
 152                 scope_fvs = unionFVInfo body_fvs rec_rhs_fvs
 153         in
 154         mapAndUnzip3Lne (varsRhs scope_fvs) pairs `thenLne` \ (rhss2, fvss, _) ->
 155         let
 156                 fvs = unionFVInfos fvss
 157         in
 158         returnLne (StgRec (binders' `zip` rhss2), fvs)
 159     )
 160
 161 \end{code}
 162
 163 \begin{code}
 164 varsRhs :: FreeVarsInfo         -- Free var info for the scope of the binding
 165         -> (Id,StgRhs)
 166         -> LneM (StgRhs, FreeVarsInfo, EscVarsSet)
 167
 168 varsRhs scope_fv_info (binder, StgRhsCon cc con args)
 169   = varsAtoms args      `thenLne` \ (args', fvs) ->
 170     returnLne (StgRhsCon cc con args', fvs, getFVSet fvs)
 171
 172 varsRhs scope_fv_info (binder, StgRhsClosure cc _ srt _ upd args body)
 173   = extendVarEnvLne [ (zapArity a, LambdaBound) | a <- args ] (
 174     do_body args body   `thenLne` \ (body2, body_fvs, body_escs) ->
 175     let
 176         set_of_args     = mkVarSet args
 177         rhs_fvs         = body_fvs  `minusFVBinders` args
 178         rhs_escs        = body_escs `minusVarSet`   set_of_args
 179         binder_info     = lookupFVInfo scope_fv_info binder
 180         upd'  | null args && isPAP body2 = ReEntrant
 181               | otherwise                = upd
 182     in
 183     returnLne (StgRhsClosure cc binder_info srt (getFVs rhs_fvs) upd'
 184                 args body2, rhs_fvs, rhs_escs)
 185     )
 186   where
 187         -- Pick out special case of application in body of thunk
 188     do_body [] (StgApp f args) = varsApp (Just upd) f args
 189     do_body _ other_body         = varsExpr other_body
 190 \end{code}
 191
 192 Detect thunks which will reduce immediately to PAPs, and make them
 193 non-updatable.  This has several advantages:
 194
 195         - the non-updatable thunk behaves exactly like the PAP,
 196
 197         - the thunk is more efficient to enter, because it is
 198           specialised to the task.
 199
 200         - we save one update frame, one stg_update_PAP, one update
 201           and lots of PAP_enters.
 202
 203         - in the case where the thunk is top-level, we save building
 204           a black hole and futhermore the thunk isn't considered to
 205           be a CAF any more, so it doesn't appear in any SRTs.
 206
 207 We do it here, because the arity information is accurate, and we need
 208 to do it before the SRT pass to save the SRT entries associated with
 209 any top-level PAPs.
 210
 211 \begin{code}
 212 isPAP (StgApp f args)
 213   = case getIdArity f of
 214            ArityExactly n -> n > n_args
 215            ArityAtLeast n -> n > n_args
 216            _              -> False
 217    where n_args = length args
 218 isPAP _ = False
 219 \end{code}
 220
 221 \begin{code}
 222 varsAtoms :: [StgArg]
 223           -> LneM ([StgArg], FreeVarsInfo)
 224         -- It's not *really* necessary to return fresh arguments,
 225         -- because the only difference is that the argument variable
 226         -- arities are correct.  But it seems safer to do so.
 227
 228 varsAtoms atoms
 229   = mapAndUnzipLne var_atom atoms       `thenLne` \ (args', fvs_lists) ->
 230     returnLne (args', unionFVInfos fvs_lists)
 231   where
 232     var_atom a@(StgConArg _) = returnLne (a, emptyFVInfo)
 233     var_atom a@(StgVarArg v)
 234       = lookupVarLne v  `thenLne` \ (v', how_bound) ->
 235         returnLne (StgVarArg v', singletonFVInfo v' how_bound stgArgOcc)
 236 \end{code}
 237
 238 %************************************************************************
 239 %*                                                                      *
 240 \subsection[expr-StgVarInfo]{Setting variable info on expressions}
 241 %*                                                                      *
 242 %************************************************************************
 243
 244 @varsExpr@ carries in a monad-ised environment, which binds each
 245 let(rec) variable (ie non top level, not imported, not lambda bound,
 246 not case-alternative bound) to:
 247         - its STG arity, and
 248         - its set of live vars.
 249 For normal variables the set of live vars is just the variable
 250 itself.  For let-no-escaped variables, the set of live vars is the set
 251 live at the moment the variable is entered.  The set is guaranteed to
 252 have no further let-no-escaped vars in it.
 253
 254 \begin{code}
 255 varsExpr :: StgExpr
 256          -> LneM (StgExpr,      -- Decorated expr
 257                   FreeVarsInfo, -- Its free vars (NB free, not live)
 258                   EscVarsSet)   -- Its escapees, a subset of its free vars;
 259                                 -- also a subset of the domain of the envt
 260                                 -- because we are only interested in the escapees
 261                                 -- for vars which might be turned into
 262                                 -- let-no-escaped ones.
 263 \end{code}
 264
 265 The second and third components can be derived in a simple bottom up pass, not
 266 dependent on any decisions about which variables will be let-no-escaped or
 267 not.  The first component, that is, the decorated expression, may then depend
 268 on these components, but it in turn is not scrutinised as the basis for any
 269 decisions.  Hence no black holes.
 270
 271 \begin{code}
 272 varsExpr (StgApp f args) = varsApp Nothing f args
 273
 274 varsExpr (StgCon con args res_ty)
 275   = getVarsLiveInCont           `thenLne` \ live_in_cont ->
 276     varsAtoms args              `thenLne` \ (args', args_fvs) ->
 277     returnLne (StgCon con args' res_ty, args_fvs, getFVSet args_fvs)
 278
 279 varsExpr (StgSCC label expr)
 280   = varsExpr expr               `thenLne` ( \ (expr2, fvs, escs) ->
 281     returnLne (StgSCC label expr2, fvs, escs) )
 282 \end{code}
 283
 284 Cases require a little more real work.
 285 \begin{code}
 286 varsExpr (StgCase scrut _ _ bndr srt alts)
 287   = getVarsLiveInCont             `thenLne` \ live_in_cont ->
 288     extendVarEnvLne [(zapArity bndr, CaseBound)] (
 289     vars_alts alts                `thenLne` \ (alts2, alts_fvs, alts_escs) ->
 290     lookupLiveVarsForSet alts_fvs `thenLne` \ alts_lvs ->
 291     let
 292         -- determine whether the default binder is dead or not
 293         bndr'= if (bndr `elementOfFVInfo` alts_fvs)
 294                   then bndr `modifyIdInfo` (setInlinePragInfo NoInlinePragInfo)
 295                   else bndr `modifyIdInfo` (setInlinePragInfo IAmDead)
 296
 297         -- don't consider the default binder as being 'live in alts',
 298         -- since this is from the point of view of the case expr, where
 299         -- the default binder is not free.
 300         live_in_alts = live_in_cont `unionVarSet`
 301                                 (alts_lvs `minusVarSet` unitVarSet bndr)
 302     in
 303         -- we tell the scrutinee that everything live in the alts
 304         -- is live in it, too.
 305     setVarsLiveInCont live_in_alts (
 306         varsExpr scrut
 307     )                              `thenLne` \ (scrut2, scrut_fvs, scrut_escs) ->
 308     lookupLiveVarsForSet scrut_fvs `thenLne` \ scrut_lvs ->
 309     let
 310         live_in_whole_case = live_in_alts `unionVarSet` scrut_lvs
 311     in
 312     returnLne (
 313       StgCase scrut2 live_in_whole_case live_in_alts bndr' srt alts2,
 314       (scrut_fvs `unionFVInfo` alts_fvs)
 315           `minusFVBinders` [bndr],
 316       (alts_escs `unionVarSet` (getFVSet scrut_fvs))
 317           `minusVarSet` unitVarSet bndr
 318
 319     ))
 320   where
 321     vars_alts (StgAlgAlts ty alts deflt)
 322       = mapAndUnzip3Lne vars_alg_alt alts
 323                         `thenLne` \ (alts2,  alts_fvs_list,  alts_escs_list) ->
 324         let
 325             alts_fvs  = unionFVInfos alts_fvs_list
 326             alts_escs = unionVarSets alts_escs_list
 327         in
 328         vars_deflt deflt `thenLne` \ (deflt2, deflt_fvs, deflt_escs) ->
 329         returnLne (
 330             StgAlgAlts ty alts2 deflt2,
 331             alts_fvs  `unionFVInfo`   deflt_fvs,
 332             alts_escs `unionVarSet` deflt_escs
 333         )
 334       where
 335         vars_alg_alt (con, binders, worthless_use_mask, rhs)
 336           = extendVarEnvLne [(zapArity b, CaseBound) | b <- binders] (
 337             varsExpr rhs        `thenLne` \ (rhs2, rhs_fvs, rhs_escs) ->
 338             let
 339                 good_use_mask = [ b `elementOfFVInfo` rhs_fvs | b <- binders ]
 340                 -- records whether each param is used in the RHS
 341             in
 342             returnLne (
 343                 (con, binders, good_use_mask, rhs2),
 344                 rhs_fvs  `minusFVBinders` binders,
 345                 rhs_escs `minusVarSet`   mkVarSet binders       -- ToDo: remove the minusVarSet;
 346                                                         -- since escs won't include
 347                                                         -- any of these binders
 348             ))
 349
 350     vars_alts (StgPrimAlts ty alts deflt)
 351       = mapAndUnzip3Lne vars_prim_alt alts
 352                         `thenLne` \ (alts2,  alts_fvs_list,  alts_escs_list) ->
 353         let
 354             alts_fvs  = unionFVInfos alts_fvs_list
 355             alts_escs = unionVarSets alts_escs_list
 356         in
 357         vars_deflt deflt `thenLne` \ (deflt2, deflt_fvs, deflt_escs) ->
 358         returnLne (
 359             StgPrimAlts ty alts2 deflt2,
 360             alts_fvs  `unionFVInfo`   deflt_fvs,
 361             alts_escs `unionVarSet` deflt_escs
 362         )
 363       where
 364         vars_prim_alt (lit, rhs)
 365           = varsExpr rhs        `thenLne` \ (rhs2, rhs_fvs, rhs_escs) ->
 366             returnLne ((lit, rhs2), rhs_fvs, rhs_escs)
 367
 368     vars_deflt StgNoDefault
 369       = returnLne (StgNoDefault, emptyFVInfo, emptyVarSet)
 370
 371     vars_deflt (StgBindDefault rhs)
 372       = varsExpr rhs    `thenLne` \ (rhs2, rhs_fvs, rhs_escs) ->
 373         returnLne ( StgBindDefault rhs2, rhs_fvs, rhs_escs )
 374 \end{code}
 375
 376 Lets not only take quite a bit of work, but this is where we convert
 377 then to let-no-escapes, if we wish.
 378
 379 (Meanwhile, we don't expect to see let-no-escapes...)
 380 \begin{code}
 381 varsExpr (StgLetNoEscape _ _ _ _) = panic "varsExpr: unexpected StgLetNoEscape"
 382
 383 varsExpr (StgLet bind body)
 384   = isSwitchSetLne {-StgDoLetNoEscapes-} `thenLne` \ want_LNEs ->
 385
 386     (fixLne (\ ~(_, _, _, no_binder_escapes) ->
 387         let
 388             non_escaping_let = want_LNEs && no_binder_escapes
 389         in
 390         vars_let non_escaping_let bind body
 391     ))                                  `thenLne` \ (new_let, fvs, escs, _) ->
 392
 393     returnLne (new_let, fvs, escs)
 394 \end{code}
 395
 396 Applications:
 397 \begin{code}
 398 varsApp :: Maybe UpdateFlag             -- Just upd <=> this application is
 399                                         -- the rhs of a thunk binding
 400                                         --      x = [...] \upd [] -> the_app
 401                                         -- with specified update flag
 402         -> Id                           -- Function
 403         -> [StgArg]             -- Arguments
 404         -> LneM (StgExpr, FreeVarsInfo, EscVarsSet)
 405
 406 varsApp maybe_thunk_body f args
 407   = getVarsLiveInCont           `thenLne` \ live_in_cont ->
 408
 409     varsAtoms args              `thenLne` \ (args', args_fvs) ->
 410
 411     lookupVarLne f              `thenLne` \ (f', how_bound) ->
 412
 413     let
 414         n_args           = length args
 415         not_letrec_bound = not (isLetrecBound how_bound)
 416         f_arity          = getIdArity f'
 417         fun_fvs          = singletonFVInfo f' how_bound fun_occ
 418
 419         fun_occ
 420           | not_letrec_bound
 421           = NoStgBinderInfo             -- Uninteresting variable
 422
 423           | otherwise                   -- Letrec bound; must have its arity
 424           = case f_arity of
 425               ArityExactly arity
 426                 | n_args == 0 -> stgFakeFunAppOcc   -- Function Application
 427                                                     -- with no arguments.
 428                                                     -- used by the lambda lifter.
 429                 | arity > n_args -> stgUnsatOcc     -- Unsaturated
 430
 431
 432                 | arity == n_args &&
 433                   maybeToBool maybe_thunk_body ->   -- Exactly saturated,
 434                                                     -- and rhs of thunk
 435                         case maybe_thunk_body of
 436                                 Just Updatable   -> stgStdHeapOcc
 437                                 Just SingleEntry -> stgNoUpdHeapOcc
 438                                 other            -> panic "varsApp"
 439
 440                 | otherwise ->  stgNormalOcc
 441                                 -- Record only that it occurs free
 442
 443         myself = unitVarSet f'
 444
 445         fun_escs | not_letrec_bound = emptyVarSet       -- Only letrec-bound escapees are interesting
 446                  | otherwise        = case f_arity of   -- Letrec bound, so must have its arity
 447                                         ArityExactly arity
 448                                           | arity == n_args -> emptyVarSet
 449                                                 -- Function doesn't escape
 450                                           | otherwise -> myself
 451                                                 -- Inexact application; it does escape
 452
 453         -- At the moment of the call:
 454
 455         --  either the function is *not* let-no-escaped, in which case
 456         --         nothing is live except live_in_cont
 457         --      or the function *is* let-no-escaped in which case the
 458         --         variables it uses are live, but still the function
 459         --         itself is not.  PS.  In this case, the function's
 460         --         live vars should already include those of the
 461         --         continuation, but it does no harm to just union the
 462         --         two regardless.
 463
 464         -- XXX not needed?
 465         -- live_at_call
 466         --   = live_in_cont `unionVarSet` case how_bound of
 467         --                            LetrecBound _ lvs -> lvs `minusVarSet` myself
 468         --                         other             -> emptyVarSet
 469     in
 470     returnLne (
 471         StgApp f' args',
 472         fun_fvs  `unionFVInfo` args_fvs,
 473         fun_escs `unionVarSet` (getFVSet args_fvs)
 474                                 -- All the free vars of the args are disqualified
 475                                 -- from being let-no-escaped.
 476     )
 477 \end{code}
 478
 479 The magic for lets:
 480 \begin{code}
 481 vars_let :: Bool                -- True <=> yes, we are let-no-escaping this let
 482          -> StgBinding  -- bindings
 483          -> StgExpr     -- body
 484          -> LneM (StgExpr,      -- new let
 485                   FreeVarsInfo, -- variables free in the whole let
 486                   EscVarsSet,   -- variables that escape from the whole let
 487                   Bool)         -- True <=> none of the binders in the bindings
 488                                 -- is among the escaping vars
 489
 490 vars_let let_no_escape bind body
 491   = fixLne (\ ~(_, _, _, rec_bind_lvs, _, rec_body_fvs, _, _) ->
 492
 493         -- Do the bindings, setting live_in_cont to empty if
 494         -- we ain't in a let-no-escape world
 495         getVarsLiveInCont               `thenLne` \ live_in_cont ->
 496         setVarsLiveInCont
 497                 (if let_no_escape then live_in_cont else emptyVarSet)
 498                 (vars_bind rec_bind_lvs rec_body_fvs bind)
 499                                         `thenLne` \ (bind2, bind_fvs, bind_escs, env_ext) ->
 500
 501         -- The live variables of this binding are the ones which are live
 502         -- by virtue of being accessible via the free vars of the binding (lvs_from_fvs)
 503         -- together with the live_in_cont ones
 504         lookupLiveVarsForSet (bind_fvs `minusFVBinders` binders)        `thenLne` \ lvs_from_fvs ->
 505         let
 506                 bind_lvs = lvs_from_fvs `unionVarSet` live_in_cont
 507         in
 508
 509         -- bind_fvs and bind_escs still include the binders of the let(rec)
 510         -- but bind_lvs does not
 511
 512         -- Do the body
 513         extendVarEnvLne env_ext (
 514                 varsExpr body                   `thenLne` \ (body2, body_fvs, body_escs) ->
 515                 lookupLiveVarsForSet body_fvs   `thenLne` \ body_lvs ->
 516
 517                 returnLne (bind2, bind_fvs, bind_escs, bind_lvs,
 518                            body2, body_fvs, body_escs, body_lvs)
 519
 520     )) `thenLne` (\ (bind2, bind_fvs, bind_escs, bind_lvs,
 521                      body2, body_fvs, body_escs, body_lvs) ->
 522
 523
 524         -- Compute the new let-expression
 525     let
 526         new_let = if let_no_escape then
 527                      -- trace "StgLetNoEscape!" (
 528                      StgLetNoEscape live_in_whole_let bind_lvs bind2 body2
 529                      -- )
 530                   else
 531                      StgLet bind2 body2
 532
 533         free_in_whole_let
 534           = (bind_fvs `unionFVInfo` body_fvs) `minusFVBinders` binders
 535
 536         live_in_whole_let
 537           = bind_lvs `unionVarSet` (body_lvs `minusVarSet` set_of_binders)
 538
 539         real_bind_escs = if let_no_escape then
 540                             bind_escs
 541                          else
 542                             getFVSet bind_fvs
 543                             -- Everything escapes which is free in the bindings
 544
 545         let_escs = (real_bind_escs `unionVarSet` body_escs) `minusVarSet` set_of_binders
 546
 547         all_escs = bind_escs `unionVarSet` body_escs    -- Still includes binders of
 548                                                 -- this let(rec)
 549
 550         no_binder_escapes = isEmptyVarSet (set_of_binders `intersectVarSet` all_escs)
 551                 -- Mustn't depend on the passed-in let_no_escape flag, since
 552                 -- no_binder_escapes is used by the caller to derive the flag!
 553     in
 554     returnLne (
 555         new_let,
 556         free_in_whole_let,
 557         let_escs,
 558         no_binder_escapes
 559     ))
 560   where
 561     set_of_binders = mkVarSet binders
 562     binders        = case bind of
 563                         StgNonRec binder rhs -> [binder]
 564                         StgRec pairs         -> map fst pairs
 565
 566     mk_binding bind_lvs (binder,rhs)
 567         = (binder `setIdArity` ArityExactly (stgArity rhs),
 568            LetrecBound  False           -- Not top level
 569                         live_vars
 570           )
 571         where
 572            live_vars = if let_no_escape then
 573                             extendVarSet bind_lvs binder
 574                        else
 575                             unitVarSet binder
 576
 577     vars_bind :: StgLiveVars
 578               -> FreeVarsInfo                   -- Free var info for body of binding
 579               -> StgBinding
 580               -> LneM (StgBinding,
 581                        FreeVarsInfo, EscVarsSet,        -- free vars; escapee vars
 582                        [(Id, HowBound)])
 583                                          -- extension to environment
 584
 585     vars_bind rec_bind_lvs rec_body_fvs (StgNonRec binder rhs)
 586       = varsRhs rec_body_fvs (binder,rhs)       `thenLne` \ (rhs2, fvs, escs) ->
 587         let
 588             env_ext_item@(binder', _) = mk_binding rec_bind_lvs (binder,rhs)
 589         in
 590         returnLne (StgNonRec binder' rhs2, fvs, escs, [env_ext_item])
 591
 592     vars_bind rec_bind_lvs rec_body_fvs (StgRec pairs)
 593       = let
 594             env_ext  = map (mk_binding rec_bind_lvs) pairs
 595             binders' = map fst env_ext
 596         in
 597         extendVarEnvLne env_ext           (
 598         fixLne (\ ~(_, rec_rhs_fvs, _, _) ->
 599                 let
 600                         rec_scope_fvs = unionFVInfo rec_body_fvs rec_rhs_fvs
 601                 in
 602                 mapAndUnzip3Lne (varsRhs rec_scope_fvs) pairs `thenLne` \ (rhss2, fvss, escss) ->
 603                 let
 604                         fvs  = unionFVInfos      fvss
 605                         escs = unionVarSets escss
 606                 in
 607                 returnLne (StgRec (binders' `zip` rhss2), fvs, escs, env_ext)
 608         ))
 609 \end{code}
 610
 611 %************************************************************************
 612 %*                                                                      *
 613 \subsection[LNE-monad]{A little monad for this let-no-escaping pass}
 614 %*                                                                      *
 615 %************************************************************************
 616
 617 There's a lot of stuff to pass around, so we use this @LneM@ monad to
 618 help.  All the stuff here is only passed {\em down}.
 619
 620 \begin{code}
 621 type LneM a =  Bool                     -- True <=> do let-no-escapes
 622             -> IdEnv (Id, HowBound)     -- Use the Id at all occurrences; it has correct
 623                                         --      arity information inside it.
 624             -> StgLiveVars              -- vars live in continuation
 625             -> a
 626
 627 data HowBound
 628   = ImportBound
 629   | CaseBound
 630   | LambdaBound
 631   | LetrecBound
 632         Bool            -- True <=> bound at top level
 633         StgLiveVars     -- Live vars... see notes below
 634
 635 isLetrecBound (LetrecBound _ _) = True
 636 isLetrecBound other             = False
 637 \end{code}
 638
 639 For a let(rec)-bound variable, x,  we record what varibles are live if
 640 x is live.  For "normal" variables that is just x alone.  If x is
 641 a let-no-escaped variable then x is represented by a code pointer and
 642 a stack pointer (well, one for each stack).  So all of the variables
 643 needed in the execution of x are live if x is, and are therefore recorded
 644 in the LetrecBound constructor; x itself *is* included.
 645
 646 The std monad functions:
 647 \begin{code}
 648 initLne :: Bool -> LneM a -> a
 649 initLne want_LNEs m = m want_LNEs emptyVarEnv emptyVarSet
 650
 651 {-# INLINE thenLne #-}
 652 {-# INLINE thenLne_ #-}
 653 {-# INLINE returnLne #-}
 654
 655 returnLne :: a -> LneM a
 656 returnLne e sw env lvs_cont = e
 657
 658 thenLne :: LneM a -> (a -> LneM b) -> LneM b
 659 thenLne m k sw env lvs_cont
 660   = case (m sw env lvs_cont) of
 661       m_result -> k m_result sw env lvs_cont
 662
 663 thenLne_ :: LneM a -> LneM b -> LneM b
 664 thenLne_ m k sw env lvs_cont
 665   = case (m sw env lvs_cont) of
 666       _ -> k sw env lvs_cont
 667
 668 mapLne  :: (a -> LneM b)   -> [a] -> LneM [b]
 669 mapLne f [] = returnLne []
 670 mapLne f (x:xs)
 671   = f x         `thenLne` \ r  ->
 672     mapLne f xs `thenLne` \ rs ->
 673     returnLne (r:rs)
 674
 675 mapAndUnzipLne  :: (a -> LneM (b,c))   -> [a] -> LneM ([b],[c])
 676
 677 mapAndUnzipLne f [] = returnLne ([],[])
 678 mapAndUnzipLne f (x:xs)
 679   = f x                 `thenLne` \ (r1,  r2)  ->
 680     mapAndUnzipLne f xs `thenLne` \ (rs1, rs2) ->
 681     returnLne (r1:rs1, r2:rs2)
 682
 683 mapAndUnzip3Lne :: (a -> LneM (b,c,d)) -> [a] -> LneM ([b],[c],[d])
 684
 685 mapAndUnzip3Lne f []    = returnLne ([],[],[])
 686 mapAndUnzip3Lne f (x:xs)
 687   = f x                  `thenLne` \ (r1,  r2,  r3)  ->
 688     mapAndUnzip3Lne f xs `thenLne` \ (rs1, rs2, rs3) ->
 689     returnLne (r1:rs1, r2:rs2, r3:rs3)
 690
 691 fixLne :: (a -> LneM a) -> LneM a
 692 fixLne expr sw env lvs_cont = result
 693   where
 694     result = expr result sw env lvs_cont
 695 --  ^^^^^^ ------ ^^^^^^
 696 \end{code}
 697
 698 Functions specific to this monad:
 699 \begin{code}
 700 isSwitchSetLne :: LneM Bool
 701 isSwitchSetLne want_LNEs env lvs_cont
 702   = want_LNEs
 703
 704 getVarsLiveInCont :: LneM StgLiveVars
 705 getVarsLiveInCont sw env lvs_cont = lvs_cont
 706
 707 setVarsLiveInCont :: StgLiveVars -> LneM a -> LneM a
 708 setVarsLiveInCont new_lvs_cont expr sw env lvs_cont
 709   = expr sw env new_lvs_cont
 710
 711 extendVarEnvLne :: [(Id, HowBound)] -> LneM a -> LneM a
 712 extendVarEnvLne ids_w_howbound expr sw env lvs_cont
 713   = expr sw (extendVarEnvList env [(id, pair) | pair@(id,_) <- ids_w_howbound]) lvs_cont
 714
 715
 716 lookupVarLne :: Id -> LneM (Id, HowBound)
 717 lookupVarLne v sw env lvs_cont
 718   = returnLne (
 719       case (lookupVarEnv env v) of
 720         Just xx -> xx
 721         Nothing -> --false:ASSERT(not (isLocallyDefined v))
 722                    (v, ImportBound)
 723     ) sw env lvs_cont
 724
 725 -- The result of lookupLiveVarsForSet, a set of live variables, is
 726 -- only ever tacked onto a decorated expression. It is never used as
 727 -- the basis of a control decision, which might give a black hole.
 728
 729 lookupLiveVarsForSet :: FreeVarsInfo -> LneM StgLiveVars
 730
 731 lookupLiveVarsForSet fvs sw env lvs_cont
 732   = returnLne (unionVarSets (map do_one (getFVs fvs)))
 733               sw env lvs_cont
 734   where
 735     do_one v
 736       = if isLocallyDefined v then
 737             case (lookupVarEnv env v) of
 738               Just (_, LetrecBound _ lvs) -> extendVarSet lvs v
 739               Just _                        -> unitVarSet v
 740               Nothing -> pprPanic "lookupVarEnv/do_one:" (ppr v)
 741         else
 742             emptyVarSet
 743 \end{code}
 744
 745
 746 %************************************************************************
 747 %*                                                                      *
 748 \subsection[Free-var info]{Free variable information}
 749 %*                                                                      *
 750 %************************************************************************
 751
 752 \begin{code}
 753 type FreeVarsInfo = IdEnv (Id, Bool, StgBinderInfo)
 754                         -- If f is mapped to NoStgBinderInfo, that means
 755                         -- that f *is* mentioned (else it wouldn't be in the
 756                         -- IdEnv at all), but only in a saturated applications.
 757                         --
 758                         -- All case/lambda-bound things are also mapped to
 759                         -- NoStgBinderInfo, since we aren't interested in their
 760                         -- occurence info.
 761                         --
 762                         -- The Bool is True <=> the Id is top level letrec bound
 763
 764 type EscVarsSet   = IdSet
 765 \end{code}
 766
 767 \begin{code}
 768 emptyFVInfo :: FreeVarsInfo
 769 emptyFVInfo = emptyVarEnv
 770
 771 singletonFVInfo :: Id -> HowBound -> StgBinderInfo -> FreeVarsInfo
 772 singletonFVInfo id ImportBound               info = emptyVarEnv
 773 singletonFVInfo id (LetrecBound top_level _) info = unitVarEnv id (id, top_level, info)
 774 singletonFVInfo id other                     info = unitVarEnv id (id, False,     info)
 775
 776 unionFVInfo :: FreeVarsInfo -> FreeVarsInfo -> FreeVarsInfo
 777 unionFVInfo fv1 fv2 = plusVarEnv_C plusFVInfo fv1 fv2
 778
 779 unionFVInfos :: [FreeVarsInfo] -> FreeVarsInfo
 780 unionFVInfos fvs = foldr unionFVInfo emptyFVInfo fvs
 781
 782 minusFVBinders :: FreeVarsInfo -> [Id] -> FreeVarsInfo
 783 minusFVBinders fv ids = fv `delVarEnvList` ids
 784
 785 elementOfFVInfo :: Id -> FreeVarsInfo -> Bool
 786 elementOfFVInfo id fvs = maybeToBool (lookupVarEnv fvs id)
 787
 788 lookupFVInfo :: FreeVarsInfo -> Id -> StgBinderInfo
 789 lookupFVInfo fvs id = case lookupVarEnv fvs id of
 790                         Nothing         -> NoStgBinderInfo
 791                         Just (_,_,info) -> info
 792
 793 getFVs :: FreeVarsInfo -> [Id]  -- Non-top-level things only
 794 getFVs fvs = [id | (id,False,_) <- rngVarEnv fvs]
 795
 796 getFVSet :: FreeVarsInfo -> IdSet
 797 getFVSet fvs = mkVarSet (getFVs fvs)
 798
 799 plusFVInfo (id1,top1,info1) (id2,top2,info2)
 800   = ASSERT (id1 == id2 && top1 == top2)
 801     (id1, top1, combineStgBinderInfo info1 info2)
 802 \end{code}
 803
 804 \begin{code}
 805 rhsArity :: StgRhs -> Arity
 806 rhsArity (StgRhsCon _ _ _)              = 0
 807 rhsArity (StgRhsClosure _ _ _ _ _ args _) = length args
 808
 809 zapArity :: Id -> Id
 810 zapArity id = id `setIdArity` UnknownArity
 811 \end{code}
 812
 813
 814