ghc/compiler/simplStg/StgVarInfo.lhs

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