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 modifyIdInfo (`setInlinePragInfo` NoInlinePragInfo) bndr
 295                   else modifyIdInfo (`setInlinePragInfo` IAmDead)          bndr
 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 `minusVarSet` unitVarSet bndr) `unionVarSet` getFVSet scrut_fvs
 317                 -- You might think we should have scrut_escs, not (getFVSet scrut_fvs),
 318                 -- but actually we can't call, and then return from, a let-no-escape thing.
 319       )
 320     )
 321   where
 322     vars_alts (StgAlgAlts ty alts deflt)
 323       = mapAndUnzip3Lne vars_alg_alt alts
 324                         `thenLne` \ (alts2,  alts_fvs_list,  alts_escs_list) ->
 325         let
 326             alts_fvs  = unionFVInfos alts_fvs_list
 327             alts_escs = unionVarSets alts_escs_list
 328         in
 329         vars_deflt deflt `thenLne` \ (deflt2, deflt_fvs, deflt_escs) ->
 330         returnLne (
 331             StgAlgAlts ty alts2 deflt2,
 332             alts_fvs  `unionFVInfo`   deflt_fvs,
 333             alts_escs `unionVarSet` deflt_escs
 334         )
 335       where
 336         vars_alg_alt (con, binders, worthless_use_mask, rhs)
 337           = extendVarEnvLne [(zapArity b, CaseBound) | b <- binders] (
 338             varsExpr rhs        `thenLne` \ (rhs2, rhs_fvs, rhs_escs) ->
 339             let
 340                 good_use_mask = [ b `elementOfFVInfo` rhs_fvs | b <- binders ]
 341                 -- records whether each param is used in the RHS
 342             in
 343             returnLne (
 344                 (con, binders, good_use_mask, rhs2),
 345                 rhs_fvs  `minusFVBinders` binders,
 346                 rhs_escs `minusVarSet`   mkVarSet binders       -- ToDo: remove the minusVarSet;
 347                                                         -- since escs won't include
 348                                                         -- any of these binders
 349             ))
 350
 351     vars_alts (StgPrimAlts ty alts deflt)
 352       = mapAndUnzip3Lne vars_prim_alt alts
 353                         `thenLne` \ (alts2,  alts_fvs_list,  alts_escs_list) ->
 354         let
 355             alts_fvs  = unionFVInfos alts_fvs_list
 356             alts_escs = unionVarSets alts_escs_list
 357         in
 358         vars_deflt deflt `thenLne` \ (deflt2, deflt_fvs, deflt_escs) ->
 359         returnLne (
 360             StgPrimAlts ty alts2 deflt2,
 361             alts_fvs  `unionFVInfo`   deflt_fvs,
 362             alts_escs `unionVarSet` deflt_escs
 363         )
 364       where
 365         vars_prim_alt (lit, rhs)
 366           = varsExpr rhs        `thenLne` \ (rhs2, rhs_fvs, rhs_escs) ->
 367             returnLne ((lit, rhs2), rhs_fvs, rhs_escs)
 368
 369     vars_deflt StgNoDefault
 370       = returnLne (StgNoDefault, emptyFVInfo, emptyVarSet)
 371
 372     vars_deflt (StgBindDefault rhs)
 373       = varsExpr rhs    `thenLne` \ (rhs2, rhs_fvs, rhs_escs) ->
 374         returnLne ( StgBindDefault rhs2, rhs_fvs, rhs_escs )
 375 \end{code}
 376
 377 Lets not only take quite a bit of work, but this is where we convert
 378 then to let-no-escapes, if we wish.
 379
 380 (Meanwhile, we don't expect to see let-no-escapes...)
 381 \begin{code}
 382 varsExpr (StgLetNoEscape _ _ _ _) = panic "varsExpr: unexpected StgLetNoEscape"
 383
 384 varsExpr (StgLet bind body)
 385   = isSwitchSetLne {-StgDoLetNoEscapes-} `thenLne` \ want_LNEs ->
 386
 387     (fixLne (\ ~(_, _, _, no_binder_escapes) ->
 388         let
 389             non_escaping_let = want_LNEs && no_binder_escapes
 390         in
 391         vars_let non_escaping_let bind body
 392     ))                                  `thenLne` \ (new_let, fvs, escs, _) ->
 393
 394     returnLne (new_let, fvs, escs)
 395 \end{code}
 396
 397 Applications:
 398 \begin{code}
 399 varsApp :: Maybe UpdateFlag             -- Just upd <=> this application is
 400                                         -- the rhs of a thunk binding
 401                                         --      x = [...] \upd [] -> the_app
 402                                         -- with specified update flag
 403         -> Id                           -- Function
 404         -> [StgArg]             -- Arguments
 405         -> LneM (StgExpr, FreeVarsInfo, EscVarsSet)
 406
 407 varsApp maybe_thunk_body f args
 408   = getVarsLiveInCont           `thenLne` \ live_in_cont ->
 409
 410     varsAtoms args              `thenLne` \ (args', args_fvs) ->
 411
 412     lookupVarLne f              `thenLne` \ (f', how_bound) ->
 413
 414     let
 415         n_args           = length args
 416         not_letrec_bound = not (isLetrecBound how_bound)
 417         f_arity          = getIdArity f'
 418         fun_fvs          = singletonFVInfo f' how_bound fun_occ
 419
 420         fun_occ
 421           | not_letrec_bound
 422           = NoStgBinderInfo             -- Uninteresting variable
 423
 424           | otherwise                   -- Letrec bound; must have its arity
 425           = case f_arity of
 426               ArityExactly arity
 427                 | n_args == 0 -> stgFakeFunAppOcc   -- Function Application
 428                                                     -- with no arguments.
 429                                                     -- used by the lambda lifter.
 430                 | arity > n_args -> stgUnsatOcc     -- Unsaturated
 431
 432
 433                 | arity == n_args &&
 434                   maybeToBool maybe_thunk_body ->   -- Exactly saturated,
 435                                                     -- and rhs of thunk
 436                         case maybe_thunk_body of
 437                                 Just Updatable   -> stgStdHeapOcc
 438                                 Just SingleEntry -> stgNoUpdHeapOcc
 439                                 other            -> panic "varsApp"
 440
 441                 | otherwise ->  stgNormalOcc
 442                                 -- Record only that it occurs free
 443
 444         myself = unitVarSet f'
 445
 446         fun_escs | not_letrec_bound = emptyVarSet       -- Only letrec-bound escapees are interesting
 447                  | otherwise        = case f_arity of   -- Letrec bound, so must have its arity
 448                                         ArityExactly arity
 449                                           | arity == n_args -> emptyVarSet
 450                                                 -- Function doesn't escape
 451                                           | otherwise -> myself
 452                                                 -- Inexact application; it does escape
 453
 454         -- At the moment of the call:
 455
 456         --  either the function is *not* let-no-escaped, in which case
 457         --         nothing is live except live_in_cont
 458         --      or the function *is* let-no-escaped in which case the
 459         --         variables it uses are live, but still the function
 460         --         itself is not.  PS.  In this case, the function's
 461         --         live vars should already include those of the
 462         --         continuation, but it does no harm to just union the
 463         --         two regardless.
 464
 465         -- XXX not needed?
 466         -- live_at_call
 467         --   = live_in_cont `unionVarSet` case how_bound of
 468         --                            LetrecBound _ lvs -> lvs `minusVarSet` myself
 469         --                         other             -> emptyVarSet
 470     in
 471     returnLne (
 472         StgApp f' args',
 473         fun_fvs  `unionFVInfo` args_fvs,
 474         fun_escs `unionVarSet` (getFVSet args_fvs)
 475                                 -- All the free vars of the args are disqualified
 476                                 -- from being let-no-escaped.
 477     )
 478 \end{code}
 479
 480 The magic for lets:
 481 \begin{code}
 482 vars_let :: Bool                -- True <=> yes, we are let-no-escaping this let
 483          -> StgBinding  -- bindings
 484          -> StgExpr     -- body
 485          -> LneM (StgExpr,      -- new let
 486                   FreeVarsInfo, -- variables free in the whole let
 487                   EscVarsSet,   -- variables that escape from the whole let
 488                   Bool)         -- True <=> none of the binders in the bindings
 489                                 -- is among the escaping vars
 490
 491 vars_let let_no_escape bind body
 492   = fixLne (\ ~(_, _, _, rec_bind_lvs, _, rec_body_fvs, _, _) ->
 493
 494         -- Do the bindings, setting live_in_cont to empty if
 495         -- we ain't in a let-no-escape world
 496         getVarsLiveInCont               `thenLne` \ live_in_cont ->
 497         setVarsLiveInCont
 498                 (if let_no_escape then live_in_cont else emptyVarSet)
 499                 (vars_bind rec_bind_lvs rec_body_fvs bind)
 500                                         `thenLne` \ (bind2, bind_fvs, bind_escs, env_ext) ->
 501
 502         -- The live variables of this binding are the ones which are live
 503         -- by virtue of being accessible via the free vars of the binding (lvs_from_fvs)
 504         -- together with the live_in_cont ones
 505         lookupLiveVarsForSet (bind_fvs `minusFVBinders` binders)        `thenLne` \ lvs_from_fvs ->
 506         let
 507                 bind_lvs = lvs_from_fvs `unionVarSet` live_in_cont
 508         in
 509
 510         -- bind_fvs and bind_escs still include the binders of the let(rec)
 511         -- but bind_lvs does not
 512
 513         -- Do the body
 514         extendVarEnvLne env_ext (
 515                 varsExpr body                   `thenLne` \ (body2, body_fvs, body_escs) ->
 516                 lookupLiveVarsForSet body_fvs   `thenLne` \ body_lvs ->
 517
 518                 returnLne (bind2, bind_fvs, bind_escs, bind_lvs,
 519                            body2, body_fvs, body_escs, body_lvs)
 520
 521     )) `thenLne` (\ (bind2, bind_fvs, bind_escs, bind_lvs,
 522                      body2, body_fvs, body_escs, body_lvs) ->
 523
 524
 525         -- Compute the new let-expression
 526     let
 527         new_let = if let_no_escape then
 528                      -- trace "StgLetNoEscape!" (
 529                      StgLetNoEscape live_in_whole_let bind_lvs bind2 body2
 530                      -- )
 531                   else
 532                      StgLet bind2 body2
 533
 534         free_in_whole_let
 535           = (bind_fvs `unionFVInfo` body_fvs) `minusFVBinders` binders
 536
 537         live_in_whole_let
 538           = bind_lvs `unionVarSet` (body_lvs `minusVarSet` set_of_binders)
 539
 540         real_bind_escs = if let_no_escape then
 541                             bind_escs
 542                          else
 543                             getFVSet bind_fvs
 544                             -- Everything escapes which is free in the bindings
 545
 546         let_escs = (real_bind_escs `unionVarSet` body_escs) `minusVarSet` set_of_binders
 547
 548         all_escs = bind_escs `unionVarSet` body_escs    -- Still includes binders of
 549                                                 -- this let(rec)
 550
 551         no_binder_escapes = isEmptyVarSet (set_of_binders `intersectVarSet` all_escs)
 552                 -- Mustn't depend on the passed-in let_no_escape flag, since
 553                 -- no_binder_escapes is used by the caller to derive the flag!
 554     in
 555     returnLne (
 556         new_let,
 557         free_in_whole_let,
 558         let_escs,
 559         no_binder_escapes
 560     ))
 561   where
 562     set_of_binders = mkVarSet binders
 563     binders        = case bind of
 564                         StgNonRec binder rhs -> [binder]
 565                         StgRec pairs         -> map fst pairs
 566
 567     mk_binding bind_lvs (binder,rhs)
 568         = (binder `setIdArity` ArityExactly (stgArity rhs),
 569            LetrecBound  False           -- Not top level
 570                         live_vars
 571           )
 572         where
 573            live_vars = if let_no_escape then
 574                             extendVarSet bind_lvs binder
 575                        else
 576                             unitVarSet binder
 577
 578     vars_bind :: StgLiveVars
 579               -> FreeVarsInfo                   -- Free var info for body of binding
 580               -> StgBinding
 581               -> LneM (StgBinding,
 582                        FreeVarsInfo, EscVarsSet,        -- free vars; escapee vars
 583                        [(Id, HowBound)])
 584                                          -- extension to environment
 585
 586     vars_bind rec_bind_lvs rec_body_fvs (StgNonRec binder rhs)
 587       = varsRhs rec_body_fvs (binder,rhs)       `thenLne` \ (rhs2, fvs, escs) ->
 588         let
 589             env_ext_item@(binder', _) = mk_binding rec_bind_lvs (binder,rhs)
 590         in
 591         returnLne (StgNonRec binder' rhs2, fvs, escs, [env_ext_item])
 592
 593     vars_bind rec_bind_lvs rec_body_fvs (StgRec pairs)
 594       = let
 595             env_ext  = map (mk_binding rec_bind_lvs) pairs
 596             binders' = map fst env_ext
 597         in
 598         extendVarEnvLne env_ext           (
 599         fixLne (\ ~(_, rec_rhs_fvs, _, _) ->
 600                 let
 601                         rec_scope_fvs = unionFVInfo rec_body_fvs rec_rhs_fvs
 602                 in
 603                 mapAndUnzip3Lne (varsRhs rec_scope_fvs) pairs `thenLne` \ (rhss2, fvss, escss) ->
 604                 let
 605                         fvs  = unionFVInfos      fvss
 606                         escs = unionVarSets escss
 607                 in
 608                 returnLne (StgRec (binders' `zip` rhss2), fvs, escs, env_ext)
 609         ))
 610 \end{code}
 611
 612 %************************************************************************
 613 %*                                                                      *
 614 \subsection[LNE-monad]{A little monad for this let-no-escaping pass}
 615 %*                                                                      *
 616 %************************************************************************
 617
 618 There's a lot of stuff to pass around, so we use this @LneM@ monad to
 619 help.  All the stuff here is only passed {\em down}.
 620
 621 \begin{code}
 622 type LneM a =  Bool                     -- True <=> do let-no-escapes
 623             -> IdEnv (Id, HowBound)     -- Use the Id at all occurrences; it has correct
 624                                         --      arity information inside it.
 625             -> StgLiveVars              -- vars live in continuation
 626             -> a
 627
 628 data HowBound
 629   = ImportBound
 630   | CaseBound
 631   | LambdaBound
 632   | LetrecBound
 633         Bool            -- True <=> bound at top level
 634         StgLiveVars     -- Live vars... see notes below
 635
 636 isLetrecBound (LetrecBound _ _) = True
 637 isLetrecBound other             = False
 638 \end{code}
 639
 640 For a let(rec)-bound variable, x,  we record what varibles are live if
 641 x is live.  For "normal" variables that is just x alone.  If x is
 642 a let-no-escaped variable then x is represented by a code pointer and
 643 a stack pointer (well, one for each stack).  So all of the variables
 644 needed in the execution of x are live if x is, and are therefore recorded
 645 in the LetrecBound constructor; x itself *is* included.
 646
 647 The std monad functions:
 648 \begin{code}
 649 initLne :: Bool -> LneM a -> a
 650 initLne want_LNEs m = m want_LNEs emptyVarEnv emptyVarSet
 651
 652 {-# INLINE thenLne #-}
 653 {-# INLINE thenLne_ #-}
 654 {-# INLINE returnLne #-}
 655
 656 returnLne :: a -> LneM a
 657 returnLne e sw env lvs_cont = e
 658
 659 thenLne :: LneM a -> (a -> LneM b) -> LneM b
 660 thenLne m k sw env lvs_cont
 661   = case (m sw env lvs_cont) of
 662       m_result -> k m_result sw env lvs_cont
 663
 664 thenLne_ :: LneM a -> LneM b -> LneM b
 665 thenLne_ m k sw env lvs_cont
 666   = case (m sw env lvs_cont) of
 667       _ -> k sw env lvs_cont
 668
 669 mapLne  :: (a -> LneM b)   -> [a] -> LneM [b]
 670 mapLne f [] = returnLne []
 671 mapLne f (x:xs)
 672   = f x         `thenLne` \ r  ->
 673     mapLne f xs `thenLne` \ rs ->
 674     returnLne (r:rs)
 675
 676 mapAndUnzipLne  :: (a -> LneM (b,c))   -> [a] -> LneM ([b],[c])
 677
 678 mapAndUnzipLne f [] = returnLne ([],[])
 679 mapAndUnzipLne f (x:xs)
 680   = f x                 `thenLne` \ (r1,  r2)  ->
 681     mapAndUnzipLne f xs `thenLne` \ (rs1, rs2) ->
 682     returnLne (r1:rs1, r2:rs2)
 683
 684 mapAndUnzip3Lne :: (a -> LneM (b,c,d)) -> [a] -> LneM ([b],[c],[d])
 685
 686 mapAndUnzip3Lne f []    = returnLne ([],[],[])
 687 mapAndUnzip3Lne f (x:xs)
 688   = f x                  `thenLne` \ (r1,  r2,  r3)  ->
 689     mapAndUnzip3Lne f xs `thenLne` \ (rs1, rs2, rs3) ->
 690     returnLne (r1:rs1, r2:rs2, r3:rs3)
 691
 692 fixLne :: (a -> LneM a) -> LneM a
 693 fixLne expr sw env lvs_cont = result
 694   where
 695     result = expr result sw env lvs_cont
 696 --  ^^^^^^ ------ ^^^^^^
 697 \end{code}
 698
 699 Functions specific to this monad:
 700 \begin{code}
 701 isSwitchSetLne :: LneM Bool
 702 isSwitchSetLne want_LNEs env lvs_cont
 703   = want_LNEs
 704
 705 getVarsLiveInCont :: LneM StgLiveVars
 706 getVarsLiveInCont sw env lvs_cont = lvs_cont
 707
 708 setVarsLiveInCont :: StgLiveVars -> LneM a -> LneM a
 709 setVarsLiveInCont new_lvs_cont expr sw env lvs_cont
 710   = expr sw env new_lvs_cont
 711
 712 extendVarEnvLne :: [(Id, HowBound)] -> LneM a -> LneM a
 713 extendVarEnvLne ids_w_howbound expr sw env lvs_cont
 714   = expr sw (extendVarEnvList env [(id, pair) | pair@(id,_) <- ids_w_howbound]) lvs_cont
 715
 716
 717 lookupVarLne :: Id -> LneM (Id, HowBound)
 718 lookupVarLne v sw env lvs_cont
 719   = returnLne (
 720       case (lookupVarEnv env v) of
 721         Just xx -> xx
 722         Nothing -> --false:ASSERT(not (isLocallyDefined v))
 723                    (v, ImportBound)
 724     ) sw env lvs_cont
 725
 726 -- The result of lookupLiveVarsForSet, a set of live variables, is
 727 -- only ever tacked onto a decorated expression. It is never used as
 728 -- the basis of a control decision, which might give a black hole.
 729
 730 lookupLiveVarsForSet :: FreeVarsInfo -> LneM StgLiveVars
 731
 732 lookupLiveVarsForSet fvs sw env lvs_cont
 733   = returnLne (unionVarSets (map do_one (getFVs fvs)))
 734               sw env lvs_cont
 735   where
 736     do_one v
 737       = if isLocallyDefined v then
 738             case (lookupVarEnv env v) of
 739               Just (_, LetrecBound _ lvs) -> extendVarSet lvs v
 740               Just _                        -> unitVarSet v
 741               Nothing -> pprPanic "lookupVarEnv/do_one:" (ppr v)
 742         else
 743             emptyVarSet
 744 \end{code}
 745
 746
 747 %************************************************************************
 748 %*                                                                      *
 749 \subsection[Free-var info]{Free variable information}
 750 %*                                                                      *
 751 %************************************************************************
 752
 753 \begin{code}
 754 type FreeVarsInfo = IdEnv (Id, Bool, StgBinderInfo)
 755                         -- If f is mapped to NoStgBinderInfo, that means
 756                         -- that f *is* mentioned (else it wouldn't be in the
 757                         -- IdEnv at all), but only in a saturated applications.
 758                         --
 759                         -- All case/lambda-bound things are also mapped to
 760                         -- NoStgBinderInfo, since we aren't interested in their
 761                         -- occurence info.
 762                         --
 763                         -- The Bool is True <=> the Id is top level letrec bound
 764
 765 type EscVarsSet   = IdSet
 766 \end{code}
 767
 768 \begin{code}
 769 emptyFVInfo :: FreeVarsInfo
 770 emptyFVInfo = emptyVarEnv
 771
 772 singletonFVInfo :: Id -> HowBound -> StgBinderInfo -> FreeVarsInfo
 773 singletonFVInfo id ImportBound               info = emptyVarEnv
 774 singletonFVInfo id (LetrecBound top_level _) info = unitVarEnv id (id, top_level, info)
 775 singletonFVInfo id other                     info = unitVarEnv id (id, False,     info)
 776
 777 unionFVInfo :: FreeVarsInfo -> FreeVarsInfo -> FreeVarsInfo
 778 unionFVInfo fv1 fv2 = plusVarEnv_C plusFVInfo fv1 fv2
 779
 780 unionFVInfos :: [FreeVarsInfo] -> FreeVarsInfo
 781 unionFVInfos fvs = foldr unionFVInfo emptyFVInfo fvs
 782
 783 minusFVBinders :: FreeVarsInfo -> [Id] -> FreeVarsInfo
 784 minusFVBinders fv ids = fv `delVarEnvList` ids
 785
 786 elementOfFVInfo :: Id -> FreeVarsInfo -> Bool
 787 elementOfFVInfo id fvs = maybeToBool (lookupVarEnv fvs id)
 788
 789 lookupFVInfo :: FreeVarsInfo -> Id -> StgBinderInfo
 790 lookupFVInfo fvs id = case lookupVarEnv fvs id of
 791                         Nothing         -> NoStgBinderInfo
 792                         Just (_,_,info) -> info
 793
 794 getFVs :: FreeVarsInfo -> [Id]  -- Non-top-level things only
 795 getFVs fvs = [id | (id,False,_) <- rngVarEnv fvs]
 796
 797 getFVSet :: FreeVarsInfo -> IdSet
 798 getFVSet fvs = mkVarSet (getFVs fvs)
 799
 800 plusFVInfo (id1,top1,info1) (id2,top2,info2)
 801   = ASSERT (id1 == id2 && top1 == top2)
 802     (id1, top1, combineStgBinderInfo info1 info2)
 803 \end{code}
 804
 805 \begin{code}
 806 rhsArity :: StgRhs -> Arity
 807 rhsArity (StgRhsCon _ _ _)              = 0
 808 rhsArity (StgRhsClosure _ _ _ _ _ args _) = length args
 809
 810 zapArity :: Id -> Id
 811 zapArity id = id `setIdArity` UnknownArity
 812 \end{code}
 813
 814
 815