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