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