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