ghc/compiler/simplCore/SetLevels.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   3 %
   4 \section{SetLevels}
   5
   6                 ***************************
   7                         Overview
   8                 ***************************
   9
  10 1. We attach binding levels to Core bindings, in preparation for floating
  11    outwards (@FloatOut@).
  12
  13 2. We also let-ify many expressions (notably case scrutinees), so they
  14    will have a fighting chance of being floated sensible.
  15
  16 3. We clone the binders of any floatable let-binding, so that when it is
  17    floated out it will be unique.  (This used to be done by the simplifier
  18    but the latter now only ensures that there's no shadowing; indeed, even
  19    that may not be true.)
  20
  21    NOTE: this can't be done using the uniqAway idea, because the variable
  22          must be unique in the whole program, not just its current scope,
  23          because two variables in different scopes may float out to the
  24          same top level place
  25
  26    NOTE: Very tiresomely, we must apply this substitution to
  27          the rules stored inside a variable too.
  28
  29    We do *not* clone top-level bindings, because some of them must not change,
  30    but we *do* clone bindings that are heading for the top level
  31
  32 4. In the expression
  33         case x of wild { p -> ...wild... }
  34    we substitute x for wild in the RHS of the case alternatives:
  35         case x of wild { p -> ...x... }
  36    This means that a sub-expression involving x is not "trapped" inside the RHS.
  37    And it's not inconvenient because we already have a substitution.
  38
  39   Note that this is EXACTLY BACKWARDS from the what the simplifier does.
  40   The simplifier tries to get rid of occurrences of x, in favour of wild,
  41   in the hope that there will only be one remaining occurrence of x, namely
  42   the scrutinee of the case, and we can inline it.
  43
  44 \begin{code}
  45 module SetLevels (
  46         setLevels,
  47
  48         Level(..), tOP_LEVEL,
  49
  50         incMinorLvl, ltMajLvl, ltLvl, isTopLvl, isInlineCtxt
  51     ) where
  52
  53 #include "HsVersions.h"
  54
  55 import CoreSyn
  56
  57 import CmdLineOpts      ( FloatOutSwitches(..) )
  58 import CoreUtils        ( exprType, exprIsTrivial, exprIsCheap, mkPiTypes )
  59 import CoreFVs          -- all of it
  60 import Subst
  61 import Id               ( Id, idType, mkSysLocal, isOneShotLambda, zapDemandIdInfo,
  62                           idSpecialisation, idWorkerInfo, setIdInfo
  63                         )
  64 import IdInfo           ( workerExists, vanillaIdInfo, )
  65 import Var              ( Var )
  66 import VarSet
  67 import VarEnv
  68 import Name             ( getOccName )
  69 import OccName          ( occNameUserString )
  70 import Type             ( isUnLiftedType, Type )
  71 import BasicTypes       ( TopLevelFlag(..) )
  72 import UniqSupply
  73 import Util             ( sortLt, isSingleton, count )
  74 import Outputable
  75 \end{code}
  76
  77 %************************************************************************
  78 %*                                                                      *
  79 \subsection{Level numbers}
  80 %*                                                                      *
  81 %************************************************************************
  82
  83 \begin{code}
  84 data Level = InlineCtxt -- A level that's used only for
  85                         -- the context parameter ctxt_lvl
  86            | Level Int  -- Level number of enclosing lambdas
  87                    Int  -- Number of big-lambda and/or case expressions between
  88                         -- here and the nearest enclosing lambda
  89 \end{code}
  90
  91 The {\em level number} on a (type-)lambda-bound variable is the
  92 nesting depth of the (type-)lambda which binds it.  The outermost lambda
  93 has level 1, so (Level 0 0) means that the variable is bound outside any lambda.
  94
  95 On an expression, it's the maximum level number of its free
  96 (type-)variables.  On a let(rec)-bound variable, it's the level of its
  97 RHS.  On a case-bound variable, it's the number of enclosing lambdas.
  98
  99 Top-level variables: level~0.  Those bound on the RHS of a top-level
 100 definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
 101 as ``subscripts'')...
 102 \begin{verbatim}
 103 a_0 = let  b_? = ...  in
 104            x_1 = ... b ... in ...
 105 \end{verbatim}
 106
 107 The main function @lvlExpr@ carries a ``context level'' (@ctxt_lvl@).
 108 That's meant to be the level number of the enclosing binder in the
 109 final (floated) program.  If the level number of a sub-expression is
 110 less than that of the context, then it might be worth let-binding the
 111 sub-expression so that it will indeed float.
 112
 113 If you can float to level @Level 0 0@ worth doing so because then your
 114 allocation becomes static instead of dynamic.  We always start with
 115 context @Level 0 0@.
 116
 117
 118 InlineCtxt
 119 ~~~~~~~~~~
 120 @InlineCtxt@ very similar to @Level 0 0@, but is used for one purpose:
 121 to say "don't float anything out of here".  That's exactly what we
 122 want for the body of an INLINE, where we don't want to float anything
 123 out at all.  See notes with lvlMFE below.
 124
 125 But, check this out:
 126
 127 -- At one time I tried the effect of not float anything out of an InlineMe,
 128 -- but it sometimes works badly.  For example, consider PrelArr.done.  It
 129 -- has the form         __inline (\d. e)
 130 -- where e doesn't mention d.  If we float this to
 131 --      __inline (let x = e in \d. x)
 132 -- things are bad.  The inliner doesn't even inline it because it doesn't look
 133 -- like a head-normal form.  So it seems a lesser evil to let things float.
 134 -- In SetLevels we do set the context to (Level 0 0) when we get to an InlineMe
 135 -- which discourages floating out.
 136
 137 So the conclusion is: don't do any floating at all inside an InlineMe.
 138 (In the above example, don't float the {x=e} out of the \d.)
 139
 140 One particular case is that of workers: we don't want to float the
 141 call to the worker outside the wrapper, otherwise the worker might get
 142 inlined into the floated expression, and an importing module won't see
 143 the worker at all.
 144
 145 \begin{code}
 146 type LevelledExpr  = TaggedExpr Level
 147 type LevelledBind  = TaggedBind Level
 148
 149 tOP_LEVEL   = Level 0 0
 150 iNLINE_CTXT = InlineCtxt
 151
 152 incMajorLvl :: Level -> Level
 153 -- For InlineCtxt we ignore any inc's; we don't want
 154 -- to do any floating at all; see notes above
 155 incMajorLvl InlineCtxt          = InlineCtxt
 156 incMajorLvl (Level major minor) = Level (major+1) 0
 157
 158 incMinorLvl :: Level -> Level
 159 incMinorLvl InlineCtxt          = InlineCtxt
 160 incMinorLvl (Level major minor) = Level major (minor+1)
 161
 162 maxLvl :: Level -> Level -> Level
 163 maxLvl InlineCtxt l2  = l2
 164 maxLvl l1  InlineCtxt = l1
 165 maxLvl l1@(Level maj1 min1) l2@(Level maj2 min2)
 166   | (maj1 > maj2) || (maj1 == maj2 && min1 > min2) = l1
 167   | otherwise                                      = l2
 168
 169 ltLvl :: Level -> Level -> Bool
 170 ltLvl any_lvl    InlineCtxt  = False
 171 ltLvl InlineCtxt (Level _ _) = True
 172 ltLvl (Level maj1 min1) (Level maj2 min2)
 173   = (maj1 < maj2) || (maj1 == maj2 && min1 < min2)
 174
 175 ltMajLvl :: Level -> Level -> Bool
 176     -- Tells if one level belongs to a difft *lambda* level to another
 177 ltMajLvl any_lvl        InlineCtxt     = False
 178 ltMajLvl InlineCtxt     (Level maj2 _) = 0 < maj2
 179 ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2
 180
 181 isTopLvl :: Level -> Bool
 182 isTopLvl (Level 0 0) = True
 183 isTopLvl other       = False
 184
 185 isInlineCtxt :: Level -> Bool
 186 isInlineCtxt InlineCtxt = True
 187 isInlineCtxt other      = False
 188
 189 instance Outputable Level where
 190   ppr InlineCtxt      = text "<INLINE>"
 191   ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]
 192
 193 instance Eq Level where
 194   InlineCtxt        == InlineCtxt        = True
 195   (Level maj1 min1) == (Level maj2 min2) = maj1==maj2 && min1==min2
 196   l1                == l2                = False
 197 \end{code}
 198
 199
 200 %************************************************************************
 201 %*                                                                      *
 202 \subsection{Main level-setting code}
 203 %*                                                                      *
 204 %************************************************************************
 205
 206 \begin{code}
 207 setLevels :: FloatOutSwitches
 208           -> [CoreBind]
 209           -> UniqSupply
 210           -> [LevelledBind]
 211
 212 setLevels float_lams binds us
 213   = initLvl us (do_them binds)
 214   where
 215     -- "do_them"'s main business is to thread the monad along
 216     -- It gives each top binding the same empty envt, because
 217     -- things unbound in the envt have level number zero implicitly
 218     do_them :: [CoreBind] -> LvlM [LevelledBind]
 219
 220     do_them [] = returnLvl []
 221     do_them (b:bs)
 222       = lvlTopBind init_env b   `thenLvl` \ (lvld_bind, _) ->
 223         do_them bs              `thenLvl` \ lvld_binds ->
 224         returnLvl (lvld_bind : lvld_binds)
 225
 226     init_env = initialEnv float_lams
 227
 228 lvlTopBind env (NonRec binder rhs)
 229   = lvlBind TopLevel tOP_LEVEL env (AnnNonRec binder (freeVars rhs))
 230                                         -- Rhs can have no free vars!
 231
 232 lvlTopBind env (Rec pairs)
 233   = lvlBind TopLevel tOP_LEVEL env (AnnRec [(b,freeVars rhs) | (b,rhs) <- pairs])
 234 \end{code}
 235
 236 %************************************************************************
 237 %*                                                                      *
 238 \subsection{Setting expression levels}
 239 %*                                                                      *
 240 %************************************************************************
 241
 242 \begin{code}
 243 lvlExpr :: Level                -- ctxt_lvl: Level of enclosing expression
 244         -> LevelEnv             -- Level of in-scope names/tyvars
 245         -> CoreExprWithFVs      -- input expression
 246         -> LvlM LevelledExpr    -- Result expression
 247 \end{code}
 248
 249 The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
 250 binder.  Here's an example
 251
 252         v = \x -> ...\y -> let r = case (..x..) of
 253                                         ..x..
 254                            in ..
 255
 256 When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
 257 the level of @r@, even though it's inside a level-2 @\y@.  It's
 258 important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
 259 don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
 260 --- because it isn't a *maximal* free expression.
 261
 262 If there were another lambda in @r@'s rhs, it would get level-2 as well.
 263
 264 \begin{code}
 265 lvlExpr _ _ (_, AnnType ty)   = returnLvl (Type ty)
 266 lvlExpr _ env (_, AnnVar v)   = returnLvl (lookupVar env v)
 267 lvlExpr _ env (_, AnnLit lit) = returnLvl (Lit lit)
 268
 269 lvlExpr ctxt_lvl env (_, AnnApp fun arg)
 270   = lvl_fun fun                         `thenLvl` \ fun' ->
 271     lvlMFE  False ctxt_lvl env arg      `thenLvl` \ arg' ->
 272     returnLvl (App fun' arg')
 273   where
 274     lvl_fun (_, AnnCase _ _ _) = lvlMFE True ctxt_lvl env fun
 275     lvl_fun other              = lvlExpr ctxt_lvl env fun
 276         -- We don't do MFE on partial applications generally,
 277         -- but we do if the function is big and hairy, like a case
 278
 279 lvlExpr ctxt_lvl env (_, AnnNote InlineMe expr)
 280 -- Don't float anything out of an InlineMe; hence the iNLINE_CTXT
 281   = lvlExpr iNLINE_CTXT env expr        `thenLvl` \ expr' ->
 282     returnLvl (Note InlineMe expr')
 283
 284 lvlExpr ctxt_lvl env (_, AnnNote note expr)
 285   = lvlExpr ctxt_lvl env expr           `thenLvl` \ expr' ->
 286     returnLvl (Note note expr')
 287
 288 -- We don't split adjacent lambdas.  That is, given
 289 --      \x y -> (x+1,y)
 290 -- we don't float to give
 291 --      \x -> let v = x+y in \y -> (v,y)
 292 -- Why not?  Because partial applications are fairly rare, and splitting
 293 -- lambdas makes them more expensive.
 294
 295 lvlExpr ctxt_lvl env expr@(_, AnnLam bndr rhs)
 296   = lvlMFE True new_lvl new_env body    `thenLvl` \ new_body ->
 297     returnLvl (mkLams new_bndrs new_body)
 298   where
 299     (bndrs, body)        = collectAnnBndrs expr
 300     (new_lvl, new_bndrs) = lvlLamBndrs ctxt_lvl bndrs
 301     new_env              = extendLvlEnv env new_bndrs
 302         -- At one time we called a special verion of collectBinders,
 303         -- which ignored coercions, because we don't want to split
 304         -- a lambda like this (\x -> coerce t (\s -> ...))
 305         -- This used to happen quite a bit in state-transformer programs,
 306         -- but not nearly so much now non-recursive newtypes are transparent.
 307         -- [See SetLevels rev 1.50 for a version with this approach.]
 308
 309 lvlExpr ctxt_lvl env (_, AnnLet bind body)
 310   = lvlBind NotTopLevel ctxt_lvl env bind       `thenLvl` \ (bind', new_env) ->
 311     lvlExpr ctxt_lvl new_env body               `thenLvl` \ body' ->
 312     returnLvl (Let bind' body')
 313
 314 lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr alts)
 315   = lvlMFE True ctxt_lvl env expr       `thenLvl` \ expr' ->
 316     let
 317         alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
 318     in
 319     mapLvl (lvl_alt alts_env) alts      `thenLvl` \ alts' ->
 320     returnLvl (Case expr' (case_bndr, incd_lvl) alts')
 321   where
 322       incd_lvl  = incMinorLvl ctxt_lvl
 323
 324       lvl_alt alts_env (con, bs, rhs)
 325         = lvlMFE True incd_lvl new_env rhs      `thenLvl` \ rhs' ->
 326           returnLvl (con, bs', rhs')
 327         where
 328           bs'     = [ (b, incd_lvl) | b <- bs ]
 329           new_env = extendLvlEnv alts_env bs'
 330 \end{code}
 331
 332 @lvlMFE@ is just like @lvlExpr@, except that it might let-bind
 333 the expression, so that it can itself be floated.
 334
 335 \begin{code}
 336 lvlMFE ::  Bool                 -- True <=> strict context [body of case or let]
 337         -> Level                -- Level of innermost enclosing lambda/tylam
 338         -> LevelEnv             -- Level of in-scope names/tyvars
 339         -> CoreExprWithFVs      -- input expression
 340         -> LvlM LevelledExpr    -- Result expression
 341
 342 lvlMFE strict_ctxt ctxt_lvl env (_, AnnType ty)
 343   = returnLvl (Type ty)
 344
 345 lvlMFE strict_ctxt ctxt_lvl env ann_expr@(fvs, _)
 346   |  isUnLiftedType ty                  -- Can't let-bind it
 347   || isInlineCtxt ctxt_lvl              -- Don't float out of an __inline__ context
 348   || exprIsTrivial expr                 -- Never float if it's trivial
 349   || not good_destination
 350   =     -- Don't float it out
 351     lvlExpr ctxt_lvl env ann_expr
 352
 353   | otherwise   -- Float it out!
 354   = lvlFloatRhs abs_vars dest_lvl env ann_expr  `thenLvl` \ expr' ->
 355     newLvlVar "lvl" abs_vars ty                 `thenLvl` \ var ->
 356     returnLvl (Let (NonRec (var,dest_lvl) expr')
 357                    (mkVarApps (Var var) abs_vars))
 358   where
 359     expr     = deAnnotate ann_expr
 360     ty       = exprType expr
 361     dest_lvl = destLevel env fvs (isFunction ann_expr)
 362     abs_vars = abstractVars dest_lvl env fvs
 363
 364         -- A decision to float entails let-binding this thing, and we only do
 365         -- that if we'll escape a value lambda, or will go to the top level.
 366     good_destination
 367         | dest_lvl `ltMajLvl` ctxt_lvl          -- Escapes a value lambda
 368         = not (exprIsCheap expr) || isTopLvl dest_lvl
 369           -- Even if it escapes a value lambda, we only
 370           -- float if it's not cheap (unless it'll get all the
 371           -- way to the top).  I've seen cases where we
 372           -- float dozens of tiny free expressions, which cost
 373           -- more to allocate than to evaluate.
 374           -- NB: exprIsCheap is also true of bottom expressions, which
 375           --     is good; we don't want to share them
 376           --
 377           -- It's only Really Bad to float a cheap expression out of a
 378           -- strict context, because that builds a thunk that otherwise
 379           -- would never be built.  So another alternative would be to
 380           -- add
 381           --    || (strict_ctxt && not (exprIsBottom expr))
 382           -- to the condition above. We should really try this out.
 383
 384         | otherwise             -- Does not escape a value lambda
 385         = isTopLvl dest_lvl     -- Only float if we are going to the top level
 386         && floatConsts env      --   and the floatConsts flag is on
 387         && not strict_ctxt      -- Don't float from a strict context
 388           -- We are keen to float something to the top level, even if it does not
 389           -- escape a lambda, because then it needs no allocation.  But it's controlled
 390           -- by a flag, because doing this too early loses opportunities for RULES
 391           -- which (needless to say) are important in some nofib programs
 392           -- (gcd is an example).
 393           --
 394           -- Beware:
 395           --    concat = /\ a -> foldr ..a.. (++) []
 396           -- was getting turned into
 397           --    concat = /\ a -> lvl a
 398           --    lvl    = /\ a -> foldr ..a.. (++) []
 399           -- which is pretty stupid.  Hence the strict_ctxt test
 400 \end{code}
 401
 402
 403 %************************************************************************
 404 %*                                                                      *
 405 \subsection{Bindings}
 406 %*                                                                      *
 407 %************************************************************************
 408
 409 The binding stuff works for top level too.
 410
 411 \begin{code}
 412 lvlBind :: TopLevelFlag         -- Used solely to decide whether to clone
 413         -> Level                -- Context level; might be Top even for bindings nested in the RHS
 414                                 -- of a top level binding
 415         -> LevelEnv
 416         -> CoreBindWithFVs
 417         -> LvlM (LevelledBind, LevelEnv)
 418
 419 lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
 420   | isInlineCtxt ctxt_lvl       -- Don't do anything inside InlineMe
 421   = lvlExpr ctxt_lvl env rhs                    `thenLvl` \ rhs' ->
 422     returnLvl (NonRec (bndr, ctxt_lvl) rhs', env)
 423
 424   | null abs_vars
 425   =     -- No type abstraction; clone existing binder
 426     lvlExpr dest_lvl env rhs                    `thenLvl` \ rhs' ->
 427     cloneVar top_lvl env bndr ctxt_lvl dest_lvl `thenLvl` \ (env', bndr') ->
 428     returnLvl (NonRec (bndr', dest_lvl) rhs', env')
 429
 430   | otherwise
 431   = -- Yes, type abstraction; create a new binder, extend substitution, etc
 432     lvlFloatRhs abs_vars dest_lvl env rhs       `thenLvl` \ rhs' ->
 433     newPolyBndrs dest_lvl env abs_vars [bndr]   `thenLvl` \ (env', [bndr']) ->
 434     returnLvl (NonRec (bndr', dest_lvl) rhs', env')
 435
 436   where
 437     bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
 438     abs_vars = abstractVars dest_lvl env bind_fvs
 439
 440     dest_lvl | isUnLiftedType (idType bndr) = destLevel env bind_fvs False `maxLvl` Level 1 0
 441              | otherwise                    = destLevel env bind_fvs (isFunction rhs)
 442         -- Hack alert!  We do have some unlifted bindings, for cheap primops, and
 443         -- it is ok to float them out; but not to the top level.  If they would otherwise
 444         -- go to the top level, we pin them inside the topmost lambda
 445 \end{code}
 446
 447
 448 \begin{code}
 449 lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
 450   | isInlineCtxt ctxt_lvl       -- Don't do anything inside InlineMe
 451   = mapLvl (lvlExpr ctxt_lvl env) rhss                  `thenLvl` \ rhss' ->
 452     returnLvl (Rec ((bndrs `zip` repeat ctxt_lvl) `zip` rhss'), env)
 453
 454   | null abs_vars
 455   = cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl    `thenLvl` \ (new_env, new_bndrs) ->
 456     mapLvl (lvlExpr ctxt_lvl new_env) rhss              `thenLvl` \ new_rhss ->
 457     returnLvl (Rec ((new_bndrs `zip` repeat dest_lvl) `zip` new_rhss), new_env)
 458
 459   | isSingleton pairs && count isId abs_vars > 1
 460   =     -- Special case for self recursion where there are
 461         -- several variables carried around: build a local loop:
 462         --      poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
 463         -- This just makes the closures a bit smaller.  If we don't do
 464         -- this, allocation rises significantly on some programs
 465         --
 466         -- We could elaborate it for the case where there are several
 467         -- mutually functions, but it's quite a bit more complicated
 468         --
 469         -- This all seems a bit ad hoc -- sigh
 470     let
 471         (bndr,rhs) = head pairs
 472         (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 473         rhs_env = extendLvlEnv env abs_vars_w_lvls
 474     in
 475     cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl   `thenLvl` \ (rhs_env', new_bndr) ->
 476     let
 477         (lam_bndrs, rhs_body)     = collectAnnBndrs rhs
 478         (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
 479         body_env                  = extendLvlEnv rhs_env' new_lam_bndrs
 480     in
 481     lvlExpr body_lvl body_env rhs_body          `thenLvl` \ new_rhs_body ->
 482     newPolyBndrs dest_lvl env abs_vars [bndr]   `thenLvl` \ (poly_env, [poly_bndr]) ->
 483     returnLvl (Rec [((poly_bndr,dest_lvl), mkLams abs_vars_w_lvls $
 484                                            mkLams new_lam_bndrs $
 485                                            Let (Rec [((new_bndr,rhs_lvl), mkLams new_lam_bndrs new_rhs_body)])
 486                                                 (mkVarApps (Var new_bndr) lam_bndrs))],
 487                poly_env)
 488
 489   | otherwise   -- Non-null abs_vars
 490   = newPolyBndrs dest_lvl env abs_vars bndrs            `thenLvl` \ (new_env, new_bndrs) ->
 491     mapLvl (lvlFloatRhs abs_vars dest_lvl new_env) rhss `thenLvl` \ new_rhss ->
 492     returnLvl (Rec ((new_bndrs `zip` repeat dest_lvl) `zip` new_rhss), new_env)
 493
 494   where
 495     (bndrs,rhss) = unzip pairs
 496
 497         -- Finding the free vars of the binding group is annoying
 498     bind_fvs        = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
 499                                     | (bndr, (rhs_fvs,_)) <- pairs])
 500                       `minusVarSet`
 501                       mkVarSet bndrs
 502
 503     dest_lvl = destLevel env bind_fvs (all isFunction rhss)
 504     abs_vars = abstractVars dest_lvl env bind_fvs
 505
 506 ----------------------------------------------------
 507 -- Three help functons for the type-abstraction case
 508
 509 lvlFloatRhs abs_vars dest_lvl env rhs
 510   = lvlExpr rhs_lvl rhs_env rhs `thenLvl` \ rhs' ->
 511     returnLvl (mkLams abs_vars_w_lvls rhs')
 512   where
 513     (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 514     rhs_env = extendLvlEnv env abs_vars_w_lvls
 515 \end{code}
 516
 517
 518 %************************************************************************
 519 %*                                                                      *
 520 \subsection{Deciding floatability}
 521 %*                                                                      *
 522 %************************************************************************
 523
 524 \begin{code}
 525 lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [(CoreBndr, Level)])
 526 -- Compute the levels for the binders of a lambda group
 527 -- The binders returned are exactly the same as the ones passed,
 528 -- but they are now paired with a level
 529 lvlLamBndrs lvl []
 530   = (lvl, [])
 531
 532 lvlLamBndrs lvl bndrs
 533   = go  (incMinorLvl lvl)
 534         False   -- Havn't bumped major level in this group
 535         [] bndrs
 536   where
 537     go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
 538         | isId bndr &&                  -- Go to the next major level if this is a value binder,
 539           not bumped_major &&           -- and we havn't already gone to the next level (one jump per group)
 540           not (isOneShotLambda bndr)    -- and it isn't a one-shot lambda
 541         = go new_lvl True ((bndr,new_lvl) : rev_lvld_bndrs) bndrs
 542
 543         | otherwise
 544         = go old_lvl bumped_major ((bndr,old_lvl) : rev_lvld_bndrs) bndrs
 545
 546         where
 547           new_lvl = incMajorLvl old_lvl
 548
 549     go old_lvl _ rev_lvld_bndrs []
 550         = (old_lvl, reverse rev_lvld_bndrs)
 551         -- a lambda like this (\x -> coerce t (\s -> ...))
 552         -- This happens quite a bit in state-transformer programs
 553 \end{code}
 554
 555 \begin{code}
 556   -- Destintion level is the max Id level of the expression
 557   -- (We'll abstract the type variables, if any.)
 558 destLevel :: LevelEnv -> VarSet -> Bool -> Level
 559 destLevel env fvs is_function
 560   |  floatLams env
 561   && is_function = tOP_LEVEL            -- Send functions to top level; see
 562                                         -- the comments with isFunction
 563   | otherwise    = maxIdLevel env fvs
 564
 565 isFunction :: CoreExprWithFVs -> Bool
 566 -- The idea here is that we want to float *functions* to
 567 -- the top level.  This saves no work, but
 568 --      (a) it can make the host function body a lot smaller,
 569 --              and hence inlinable.
 570 --      (b) it can also save allocation when the function is recursive:
 571 --          h = \x -> letrec f = \y -> ...f...y...x...
 572 --                    in f x
 573 --     becomes
 574 --          f = \x y -> ...(f x)...y...x...
 575 --          h = \x -> f x x
 576 --     No allocation for f now.
 577 -- We may only want to do this if there are sufficiently few free
 578 -- variables.  We certainly only want to do it for values, and not for
 579 -- constructors.  So the simple thing is just to look for lambdas
 580 isFunction (_, AnnLam b e) | isId b    = True
 581                            | otherwise = isFunction e
 582 isFunction (_, AnnNote n e)            = isFunction e
 583 isFunction other                       = False
 584 \end{code}
 585
 586
 587 %************************************************************************
 588 %*                                                                      *
 589 \subsection{Free-To-Level Monad}
 590 %*                                                                      *
 591 %************************************************************************
 592
 593 \begin{code}
 594 type LevelEnv = (FloatOutSwitches,
 595                  VarEnv Level,                  -- Domain is *post-cloned* TyVars and Ids
 596                  Subst,                         -- Domain is pre-cloned Ids; tracks the in-scope set
 597                                                 --      so that subtitution is capture-avoiding
 598                  IdEnv ([Var], LevelledExpr))   -- Domain is pre-cloned Ids
 599         -- We clone let-bound variables so that they are still
 600         -- distinct when floated out; hence the SubstEnv/IdEnv.
 601         -- (see point 3 of the module overview comment).
 602         -- We also use these envs when making a variable polymorphic
 603         -- because we want to float it out past a big lambda.
 604         --
 605         -- The SubstEnv and IdEnv always implement the same mapping, but the
 606         -- SubstEnv maps to CoreExpr and the IdEnv to LevelledExpr
 607         -- Since the range is always a variable or type application,
 608         -- there is never any difference between the two, but sadly
 609         -- the types differ.  The SubstEnv is used when substituting in
 610         -- a variable's IdInfo; the IdEnv when we find a Var.
 611         --
 612         -- In addition the IdEnv records a list of tyvars free in the
 613         -- type application, just so we don't have to call freeVars on
 614         -- the type application repeatedly.
 615         --
 616         -- The domain of the both envs is *pre-cloned* Ids, though
 617         --
 618         -- The domain of the VarEnv Level is the *post-cloned* Ids
 619
 620 initialEnv :: FloatOutSwitches -> LevelEnv
 621 initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
 622
 623 floatLams :: LevelEnv -> Bool
 624 floatLams (FloatOutSw float_lams _, _, _, _) = float_lams
 625
 626 floatConsts :: LevelEnv -> Bool
 627 floatConsts (FloatOutSw _ float_consts, _, _, _) = float_consts
 628
 629 extendLvlEnv :: LevelEnv -> [(Var,Level)] -> LevelEnv
 630 -- Used when *not* cloning
 631 extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
 632   = (float_lams,
 633      foldl add_lvl lvl_env prs,
 634      foldl del_subst subst prs,
 635      foldl del_id id_env prs)
 636   where
 637     add_lvl   env (v,l) = extendVarEnv env v l
 638     del_subst env (v,_) = extendInScope env v
 639     del_id    env (v,_) = delVarEnv env v
 640   -- We must remove any clone for this variable name in case of
 641   -- shadowing.  This bit me in the following case
 642   -- (in nofib/real/gg/Spark.hs):
 643   --
 644   --   case ds of wild {
 645   --     ... -> case e of wild {
 646   --              ... -> ... wild ...
 647   --            }
 648   --   }
 649   --
 650   -- The inside occurrence of @wild@ was being replaced with @ds@,
 651   -- incorrectly, because the SubstEnv was still lying around.  Ouch!
 652   -- KSW 2000-07.
 653
 654 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
 655 -- (see point 4 of the module overview comment)
 656 extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
 657   = (float_lams,
 658      extendVarEnv lvl_env case_bndr lvl,
 659      extendSubst subst case_bndr (DoneEx (Var scrut_var)),
 660      extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
 661
 662 extendCaseBndrLvlEnv env scrut case_bndr lvl
 663   = extendLvlEnv          env [(case_bndr,lvl)]
 664
 665 extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
 666   = (float_lams,
 667      foldl add_lvl   lvl_env bndr_pairs,
 668      foldl add_subst subst   bndr_pairs,
 669      foldl add_id    id_env  bndr_pairs)
 670   where
 671      add_lvl   env (v,v') = extendVarEnv env v' dest_lvl
 672      add_subst env (v,v') = extendSubst  env v (DoneEx (mkVarApps (Var v') abs_vars))
 673      add_id    env (v,v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
 674
 675 extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
 676   = (float_lams,
 677      foldl add_lvl   lvl_env bndr_pairs,
 678      new_subst,
 679      foldl add_id    id_env  bndr_pairs)
 680   where
 681      add_lvl   env (v,v') = extendVarEnv env v' lvl
 682      add_id    env (v,v') = extendVarEnv env v ([v'], Var v')
 683
 684
 685 maxIdLevel :: LevelEnv -> VarSet -> Level
 686 maxIdLevel (_, lvl_env,_,id_env) var_set
 687   = foldVarSet max_in tOP_LEVEL var_set
 688   where
 689     max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
 690                                                 Just (abs_vars, _) -> abs_vars
 691                                                 Nothing            -> [in_var])
 692
 693     max_out out_var lvl
 694         | isId out_var = case lookupVarEnv lvl_env out_var of
 695                                 Just lvl' -> maxLvl lvl' lvl
 696                                 Nothing   -> lvl
 697         | otherwise    = lvl    -- Ignore tyvars in *maxIdLevel*
 698
 699 lookupVar :: LevelEnv -> Id -> LevelledExpr
 700 lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
 701                                        Just (_, expr) -> expr
 702                                        other          -> Var v
 703
 704 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
 705         -- Find the variables in fvs, free vars of the target expresion,
 706         -- whose level is greater than the destination level
 707         -- These are the ones we are going to abstract out
 708 abstractVars dest_lvl env fvs
 709   = uniq (sortLt lt [var | fv <- varSetElems fvs, var <- absVarsOf dest_lvl env fv])
 710   where
 711         -- Sort the variables so we don't get
 712         -- mixed-up tyvars and Ids; it's just messy
 713     v1 `lt` v2 = case (isId v1, isId v2) of
 714                    (True, False) -> False
 715                    (False, True) -> True
 716                    other         -> v1 < v2     -- Same family
 717
 718     uniq :: [Var] -> [Var]
 719         -- Remove adjacent duplicates; the sort will have brought them together
 720     uniq (v1:v2:vs) | v1 == v2  = uniq (v2:vs)
 721                     | otherwise = v1 : uniq (v2:vs)
 722     uniq vs = vs
 723
 724 absVarsOf :: Level -> LevelEnv -> Var -> [Var]
 725         -- If f is free in the expression, and f maps to poly_f a b c in the
 726         -- current substitution, then we must report a b c as candidate type
 727         -- variables
 728 absVarsOf dest_lvl (_, lvl_env, _, id_env) v
 729   | isId v
 730   = [zap av2 | av1 <- lookup_avs v, av2 <- add_tyvars av1, abstract_me av2]
 731
 732   | otherwise
 733   = if abstract_me v then [v] else []
 734
 735   where
 736     abstract_me v = case lookupVarEnv lvl_env v of
 737                         Just lvl -> dest_lvl `ltLvl` lvl
 738                         Nothing  -> False
 739
 740     lookup_avs v = case lookupVarEnv id_env v of
 741                         Just (abs_vars, _) -> abs_vars
 742                         Nothing            -> [v]
 743
 744     add_tyvars v | isId v    = v : varSetElems (idFreeTyVars v)
 745                  | otherwise = [v]
 746
 747         -- We are going to lambda-abstract, so nuke any IdInfo,
 748         -- and add the tyvars of the Id (if necessary)
 749     zap v | isId v = WARN( workerExists (idWorkerInfo v) ||
 750                            not (isEmptyCoreRules (idSpecialisation v)),
 751                            text "absVarsOf: discarding info on" <+> ppr v )
 752                      setIdInfo v vanillaIdInfo
 753           | otherwise = v
 754 \end{code}
 755
 756 \begin{code}
 757 type LvlM result = UniqSM result
 758
 759 initLvl         = initUs_
 760 thenLvl         = thenUs
 761 returnLvl       = returnUs
 762 mapLvl          = mapUs
 763 \end{code}
 764
 765 \begin{code}
 766 newPolyBndrs dest_lvl env abs_vars bndrs
 767   = getUniquesUs                `thenLvl` \ uniqs ->
 768     let
 769         new_bndrs = zipWith mk_poly_bndr bndrs uniqs
 770     in
 771     returnLvl (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
 772   where
 773     mk_poly_bndr bndr uniq = mkSysLocal (_PK_ str) uniq poly_ty
 774                            where
 775                              str     = "poly_" ++ occNameUserString (getOccName bndr)
 776                              poly_ty = mkPiTypes abs_vars (idType bndr)
 777
 778
 779 newLvlVar :: String
 780           -> [CoreBndr] -> Type         -- Abstract wrt these bndrs
 781           -> LvlM Id
 782 newLvlVar str vars body_ty
 783   = getUniqueUs `thenLvl` \ uniq ->
 784     returnUs (mkSysLocal (_PK_ str) uniq (mkPiTypes vars body_ty))
 785
 786 -- The deeply tiresome thing is that we have to apply the substitution
 787 -- to the rules inside each Id.  Grr.  But it matters.
 788
 789 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
 790 cloneVar TopLevel env v ctxt_lvl dest_lvl
 791   = returnUs (env, v)   -- Don't clone top level things
 792 cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
 793   = ASSERT( isId v )
 794     getUs       `thenLvl` \ us ->
 795     let
 796       (subst', v1) = substAndCloneId subst us v
 797       v2           = zap_demand ctxt_lvl dest_lvl v1
 798       env'         = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
 799     in
 800     returnUs (env', v2)
 801
 802 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
 803 cloneRecVars TopLevel env vs ctxt_lvl dest_lvl
 804   = returnUs (env, vs)  -- Don't clone top level things
 805 cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
 806   = ASSERT( all isId vs )
 807     getUs                       `thenLvl` \ us ->
 808     let
 809       (subst', vs1) = substAndCloneRecIds subst us vs
 810       vs2           = map (zap_demand ctxt_lvl dest_lvl) vs1
 811       env'          = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
 812     in
 813     returnUs (env', vs2)
 814
 815         -- VERY IMPORTANT: we must zap the demand info
 816         -- if the thing is going to float out past a lambda
 817 zap_demand dest_lvl ctxt_lvl id
 818   | ctxt_lvl == dest_lvl = id                   -- Stays put
 819   | otherwise            = zapDemandIdInfo id   -- Floats out
 820 \end{code}
 821