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