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