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, exprIsBottom, 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   || not good_destination
 349   || exprIsTrivial expr                 -- Is trivial
 350   || (strict_ctxt && exprIsBottom expr) -- Strict context and is bottom
 351                                         --  e.g. \x -> error "foo"
 352                                         -- No gain from floating this
 353   =     -- Don't float it out
 354     lvlExpr ctxt_lvl env ann_expr
 355
 356   | otherwise   -- Float it out!
 357   = lvlFloatRhs abs_vars dest_lvl env ann_expr  `thenLvl` \ expr' ->
 358     newLvlVar "lvl" abs_vars ty                 `thenLvl` \ var ->
 359     returnLvl (Let (NonRec (var,dest_lvl) expr')
 360                    (mkVarApps (Var var) abs_vars))
 361   where
 362     expr     = deAnnotate ann_expr
 363     ty       = exprType expr
 364     dest_lvl = destLevel env fvs (isFunction ann_expr)
 365     abs_vars = abstractVars dest_lvl env fvs
 366
 367     good_destination =  dest_lvl `ltMajLvl` ctxt_lvl    -- Escapes a value lambda
 368                      || (isTopLvl dest_lvl              -- Goes to the top
 369                          && floatConsts env
 370                          && not strict_ctxt)            --   or from a strict context
 371         -- A decision to float entails let-binding this thing, and we only do
 372         -- that if we'll escape a value lambda, or will go to the top level.
 373         --
 374         -- Beware:
 375         --      concat = /\ a -> foldr ..a.. (++) []
 376         -- was getting turned into
 377         --      concat = /\ a -> lvl a
 378         --      lvl    = /\ a -> foldr ..a.. (++) []
 379         -- which is pretty stupid.  Hence the strict_ctxt test
 380         --
 381         -- We are keen to float something to the top level, even if it does not
 382         -- escape a lambda, because then it needs no allocation.  But it's controlled
 383         -- by a flag, because doing this too early loses opportunities for RULES
 384         -- which (needless to say) are important in some nofib programs
 385         -- (gcd is an example).
 386 \end{code}
 387
 388
 389 %************************************************************************
 390 %*                                                                      *
 391 \subsection{Bindings}
 392 %*                                                                      *
 393 %************************************************************************
 394
 395 The binding stuff works for top level too.
 396
 397 \begin{code}
 398 lvlBind :: TopLevelFlag         -- Used solely to decide whether to clone
 399         -> Level                -- Context level; might be Top even for bindings nested in the RHS
 400                                 -- of a top level binding
 401         -> LevelEnv
 402         -> CoreBindWithFVs
 403         -> LvlM (LevelledBind, LevelEnv)
 404
 405 lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
 406   | isInlineCtxt ctxt_lvl       -- Don't do anything inside InlineMe
 407   = lvlExpr ctxt_lvl env rhs                    `thenLvl` \ rhs' ->
 408     returnLvl (NonRec (bndr, ctxt_lvl) rhs', env)
 409
 410   | null abs_vars
 411   =     -- No type abstraction; clone existing binder
 412     lvlExpr dest_lvl env rhs                    `thenLvl` \ rhs' ->
 413     cloneVar top_lvl env bndr ctxt_lvl dest_lvl `thenLvl` \ (env', bndr') ->
 414     returnLvl (NonRec (bndr', dest_lvl) rhs', env')
 415
 416   | otherwise
 417   = -- Yes, type abstraction; create a new binder, extend substitution, etc
 418     lvlFloatRhs abs_vars dest_lvl env rhs       `thenLvl` \ rhs' ->
 419     newPolyBndrs dest_lvl env abs_vars [bndr]   `thenLvl` \ (env', [bndr']) ->
 420     returnLvl (NonRec (bndr', dest_lvl) rhs', env')
 421
 422   where
 423     bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
 424     abs_vars = abstractVars dest_lvl env bind_fvs
 425
 426     dest_lvl | isUnLiftedType (idType bndr) = destLevel env bind_fvs False `maxLvl` Level 1 0
 427              | otherwise                    = destLevel env bind_fvs (isFunction rhs)
 428         -- Hack alert!  We do have some unlifted bindings, for cheap primops, and
 429         -- it is ok to float them out; but not to the top level.  If they would otherwise
 430         -- go to the top level, we pin them inside the topmost lambda
 431 \end{code}
 432
 433
 434 \begin{code}
 435 lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
 436   | isInlineCtxt ctxt_lvl       -- Don't do anything inside InlineMe
 437   = mapLvl (lvlExpr ctxt_lvl env) rhss                  `thenLvl` \ rhss' ->
 438     returnLvl (Rec ((bndrs `zip` repeat ctxt_lvl) `zip` rhss'), env)
 439
 440   | null abs_vars
 441   = cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl    `thenLvl` \ (new_env, new_bndrs) ->
 442     mapLvl (lvlExpr ctxt_lvl new_env) rhss              `thenLvl` \ new_rhss ->
 443     returnLvl (Rec ((new_bndrs `zip` repeat dest_lvl) `zip` new_rhss), new_env)
 444
 445   | isSingleton pairs && count isId abs_vars > 1
 446   =     -- Special case for self recursion where there are
 447         -- several variables carried around: build a local loop:
 448         --      poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
 449         -- This just makes the closures a bit smaller.  If we don't do
 450         -- this, allocation rises significantly on some programs
 451         --
 452         -- We could elaborate it for the case where there are several
 453         -- mutually functions, but it's quite a bit more complicated
 454         --
 455         -- This all seems a bit ad hoc -- sigh
 456     let
 457         (bndr,rhs) = head pairs
 458         (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 459         rhs_env = extendLvlEnv env abs_vars_w_lvls
 460     in
 461     cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl   `thenLvl` \ (rhs_env', new_bndr) ->
 462     let
 463         (lam_bndrs, rhs_body)     = collectAnnBndrs rhs
 464         (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
 465         body_env                  = extendLvlEnv rhs_env' new_lam_bndrs
 466     in
 467     lvlExpr body_lvl body_env rhs_body          `thenLvl` \ new_rhs_body ->
 468     newPolyBndrs dest_lvl env abs_vars [bndr]   `thenLvl` \ (poly_env, [poly_bndr]) ->
 469     returnLvl (Rec [((poly_bndr,dest_lvl), mkLams abs_vars_w_lvls $
 470                                            mkLams new_lam_bndrs $
 471                                            Let (Rec [((new_bndr,rhs_lvl), mkLams new_lam_bndrs new_rhs_body)])
 472                                                 (mkVarApps (Var new_bndr) lam_bndrs))],
 473                poly_env)
 474
 475   | otherwise   -- Non-null abs_vars
 476   = newPolyBndrs dest_lvl env abs_vars bndrs            `thenLvl` \ (new_env, new_bndrs) ->
 477     mapLvl (lvlFloatRhs abs_vars dest_lvl new_env) rhss `thenLvl` \ new_rhss ->
 478     returnLvl (Rec ((new_bndrs `zip` repeat dest_lvl) `zip` new_rhss), new_env)
 479
 480   where
 481     (bndrs,rhss) = unzip pairs
 482
 483         -- Finding the free vars of the binding group is annoying
 484     bind_fvs        = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
 485                                     | (bndr, (rhs_fvs,_)) <- pairs])
 486                       `minusVarSet`
 487                       mkVarSet bndrs
 488
 489     dest_lvl = destLevel env bind_fvs (all isFunction rhss)
 490     abs_vars = abstractVars dest_lvl env bind_fvs
 491
 492 ----------------------------------------------------
 493 -- Three help functons for the type-abstraction case
 494
 495 lvlFloatRhs abs_vars dest_lvl env rhs
 496   = lvlExpr rhs_lvl rhs_env rhs `thenLvl` \ rhs' ->
 497     returnLvl (mkLams abs_vars_w_lvls rhs')
 498   where
 499     (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 500     rhs_env = extendLvlEnv env abs_vars_w_lvls
 501 \end{code}
 502
 503
 504 %************************************************************************
 505 %*                                                                      *
 506 \subsection{Deciding floatability}
 507 %*                                                                      *
 508 %************************************************************************
 509
 510 \begin{code}
 511 lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [(CoreBndr, Level)])
 512 -- Compute the levels for the binders of a lambda group
 513 -- The binders returned are exactly the same as the ones passed,
 514 -- but they are now paired with a level
 515 lvlLamBndrs lvl []
 516   = (lvl, [])
 517
 518 lvlLamBndrs lvl bndrs
 519   = go  (incMinorLvl lvl)
 520         False   -- Havn't bumped major level in this group
 521         [] bndrs
 522   where
 523     go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
 524         | isId bndr &&                  -- Go to the next major level if this is a value binder,
 525           not bumped_major &&           -- and we havn't already gone to the next level (one jump per group)
 526           not (isOneShotLambda bndr)    -- and it isn't a one-shot lambda
 527         = go new_lvl True ((bndr,new_lvl) : rev_lvld_bndrs) bndrs
 528
 529         | otherwise
 530         = go old_lvl bumped_major ((bndr,old_lvl) : rev_lvld_bndrs) bndrs
 531
 532         where
 533           new_lvl = incMajorLvl old_lvl
 534
 535     go old_lvl _ rev_lvld_bndrs []
 536         = (old_lvl, reverse rev_lvld_bndrs)
 537         -- a lambda like this (\x -> coerce t (\s -> ...))
 538         -- This happens quite a bit in state-transformer programs
 539 \end{code}
 540
 541 \begin{code}
 542   -- Destintion level is the max Id level of the expression
 543   -- (We'll abstract the type variables, if any.)
 544 destLevel :: LevelEnv -> VarSet -> Bool -> Level
 545 destLevel env fvs is_function
 546   |  floatLams env
 547   && is_function = tOP_LEVEL            -- Send functions to top level; see
 548                                         -- the comments with isFunction
 549   | otherwise    = maxIdLevel env fvs
 550
 551 isFunction :: CoreExprWithFVs -> Bool
 552 -- The idea here is that we want to float *functions* to
 553 -- the top level.  This saves no work, but
 554 --      (a) it can make the host function body a lot smaller,
 555 --              and hence inlinable.
 556 --      (b) it can also save allocation when the function is recursive:
 557 --          h = \x -> letrec f = \y -> ...f...y...x...
 558 --                    in f x
 559 --     becomes
 560 --          f = \x y -> ...(f x)...y...x...
 561 --          h = \x -> f x x
 562 --     No allocation for f now.
 563 -- We may only want to do this if there are sufficiently few free
 564 -- variables.  We certainly only want to do it for values, and not for
 565 -- constructors.  So the simple thing is just to look for lambdas
 566 isFunction (_, AnnLam b e) | isId b    = True
 567                            | otherwise = isFunction e
 568 isFunction (_, AnnNote n e)            = isFunction e
 569 isFunction other                       = False
 570 \end{code}
 571
 572
 573 %************************************************************************
 574 %*                                                                      *
 575 \subsection{Free-To-Level Monad}
 576 %*                                                                      *
 577 %************************************************************************
 578
 579 \begin{code}
 580 type LevelEnv = (FloatOutSwitches,
 581                  VarEnv Level,                  -- Domain is *post-cloned* TyVars and Ids
 582                  Subst,                         -- Domain is pre-cloned Ids; tracks the in-scope set
 583                                                 --      so that subtitution is capture-avoiding
 584                  IdEnv ([Var], LevelledExpr))   -- Domain is pre-cloned Ids
 585         -- We clone let-bound variables so that they are still
 586         -- distinct when floated out; hence the SubstEnv/IdEnv.
 587         -- (see point 3 of the module overview comment).
 588         -- We also use these envs when making a variable polymorphic
 589         -- because we want to float it out past a big lambda.
 590         --
 591         -- The SubstEnv and IdEnv always implement the same mapping, but the
 592         -- SubstEnv maps to CoreExpr and the IdEnv to LevelledExpr
 593         -- Since the range is always a variable or type application,
 594         -- there is never any difference between the two, but sadly
 595         -- the types differ.  The SubstEnv is used when substituting in
 596         -- a variable's IdInfo; the IdEnv when we find a Var.
 597         --
 598         -- In addition the IdEnv records a list of tyvars free in the
 599         -- type application, just so we don't have to call freeVars on
 600         -- the type application repeatedly.
 601         --
 602         -- The domain of the both envs is *pre-cloned* Ids, though
 603         --
 604         -- The domain of the VarEnv Level is the *post-cloned* Ids
 605
 606 initialEnv :: FloatOutSwitches -> LevelEnv
 607 initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
 608
 609 floatLams :: LevelEnv -> Bool
 610 floatLams (FloatOutSw float_lams _, _, _, _) = float_lams
 611
 612 floatConsts :: LevelEnv -> Bool
 613 floatConsts (FloatOutSw _ float_consts, _, _, _) = float_consts
 614
 615 extendLvlEnv :: LevelEnv -> [(Var,Level)] -> LevelEnv
 616 -- Used when *not* cloning
 617 extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
 618   = (float_lams,
 619      foldl add_lvl lvl_env prs,
 620      foldl del_subst subst prs,
 621      foldl del_id id_env prs)
 622   where
 623     add_lvl   env (v,l) = extendVarEnv env v l
 624     del_subst env (v,_) = extendInScope env v
 625     del_id    env (v,_) = delVarEnv env v
 626   -- We must remove any clone for this variable name in case of
 627   -- shadowing.  This bit me in the following case
 628   -- (in nofib/real/gg/Spark.hs):
 629   --
 630   --   case ds of wild {
 631   --     ... -> case e of wild {
 632   --              ... -> ... wild ...
 633   --            }
 634   --   }
 635   --
 636   -- The inside occurrence of @wild@ was being replaced with @ds@,
 637   -- incorrectly, because the SubstEnv was still lying around.  Ouch!
 638   -- KSW 2000-07.
 639
 640 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
 641 -- (see point 4 of the module overview comment)
 642 extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
 643   = (float_lams,
 644      extendVarEnv lvl_env case_bndr lvl,
 645      extendSubst subst case_bndr (DoneEx (Var scrut_var)),
 646      extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
 647
 648 extendCaseBndrLvlEnv env scrut case_bndr lvl
 649   = extendLvlEnv          env [(case_bndr,lvl)]
 650
 651 extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
 652   = (float_lams,
 653      foldl add_lvl   lvl_env bndr_pairs,
 654      foldl add_subst subst   bndr_pairs,
 655      foldl add_id    id_env  bndr_pairs)
 656   where
 657      add_lvl   env (v,v') = extendVarEnv env v' dest_lvl
 658      add_subst env (v,v') = extendSubst  env v (DoneEx (mkVarApps (Var v') abs_vars))
 659      add_id    env (v,v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
 660
 661 extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
 662   = (float_lams,
 663      foldl add_lvl   lvl_env bndr_pairs,
 664      new_subst,
 665      foldl add_id    id_env  bndr_pairs)
 666   where
 667      add_lvl   env (v,v') = extendVarEnv env v' lvl
 668      add_id    env (v,v') = extendVarEnv env v ([v'], Var v')
 669
 670
 671 maxIdLevel :: LevelEnv -> VarSet -> Level
 672 maxIdLevel (_, lvl_env,_,id_env) var_set
 673   = foldVarSet max_in tOP_LEVEL var_set
 674   where
 675     max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
 676                                                 Just (abs_vars, _) -> abs_vars
 677                                                 Nothing            -> [in_var])
 678
 679     max_out out_var lvl
 680         | isId out_var = case lookupVarEnv lvl_env out_var of
 681                                 Just lvl' -> maxLvl lvl' lvl
 682                                 Nothing   -> lvl
 683         | otherwise    = lvl    -- Ignore tyvars in *maxIdLevel*
 684
 685 lookupVar :: LevelEnv -> Id -> LevelledExpr
 686 lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
 687                                        Just (_, expr) -> expr
 688                                        other          -> Var v
 689
 690 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
 691         -- Find the variables in fvs, free vars of the target expresion,
 692         -- whose level is greater than the destination level
 693         -- These are the ones we are going to abstract out
 694 abstractVars dest_lvl env fvs
 695   = uniq (sortLt lt [var | fv <- varSetElems fvs, var <- absVarsOf dest_lvl env fv])
 696   where
 697         -- Sort the variables so we don't get
 698         -- mixed-up tyvars and Ids; it's just messy
 699     v1 `lt` v2 = case (isId v1, isId v2) of
 700                    (True, False) -> False
 701                    (False, True) -> True
 702                    other         -> v1 < v2     -- Same family
 703
 704     uniq :: [Var] -> [Var]
 705         -- Remove adjacent duplicates; the sort will have brought them together
 706     uniq (v1:v2:vs) | v1 == v2  = uniq (v2:vs)
 707                     | otherwise = v1 : uniq (v2:vs)
 708     uniq vs = vs
 709
 710 absVarsOf :: Level -> LevelEnv -> Var -> [Var]
 711         -- If f is free in the expression, and f maps to poly_f a b c in the
 712         -- current substitution, then we must report a b c as candidate type
 713         -- variables
 714 absVarsOf dest_lvl (_, lvl_env, _, id_env) v
 715   | isId v
 716   = [zap av2 | av1 <- lookup_avs v, av2 <- add_tyvars av1, abstract_me av2]
 717
 718   | otherwise
 719   = if abstract_me v then [v] else []
 720
 721   where
 722     abstract_me v = case lookupVarEnv lvl_env v of
 723                         Just lvl -> dest_lvl `ltLvl` lvl
 724                         Nothing  -> False
 725
 726     lookup_avs v = case lookupVarEnv id_env v of
 727                         Just (abs_vars, _) -> abs_vars
 728                         Nothing            -> [v]
 729
 730     add_tyvars v | isId v    = v : varSetElems (idFreeTyVars v)
 731                  | otherwise = [v]
 732
 733         -- We are going to lambda-abstract, so nuke any IdInfo,
 734         -- and add the tyvars of the Id (if necessary)
 735     zap v | isId v = WARN( workerExists (idWorkerInfo v) ||
 736                            not (isEmptyCoreRules (idSpecialisation v)),
 737                            text "absVarsOf: discarding info on" <+> ppr v )
 738                      setIdInfo v vanillaIdInfo
 739           | otherwise = v
 740 \end{code}
 741
 742 \begin{code}
 743 type LvlM result = UniqSM result
 744
 745 initLvl         = initUs_
 746 thenLvl         = thenUs
 747 returnLvl       = returnUs
 748 mapLvl          = mapUs
 749 \end{code}
 750
 751 \begin{code}
 752 newPolyBndrs dest_lvl env abs_vars bndrs
 753   = getUniquesUs                `thenLvl` \ uniqs ->
 754     let
 755         new_bndrs = zipWith mk_poly_bndr bndrs uniqs
 756     in
 757     returnLvl (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
 758   where
 759     mk_poly_bndr bndr uniq = mkSysLocal (_PK_ str) uniq poly_ty
 760                            where
 761                              str     = "poly_" ++ occNameUserString (getOccName bndr)
 762                              poly_ty = mkPiTypes abs_vars (idType bndr)
 763
 764
 765 newLvlVar :: String
 766           -> [CoreBndr] -> Type         -- Abstract wrt these bndrs
 767           -> LvlM Id
 768 newLvlVar str vars body_ty
 769   = getUniqueUs `thenLvl` \ uniq ->
 770     returnUs (mkSysLocal (_PK_ str) uniq (mkPiTypes vars body_ty))
 771
 772 -- The deeply tiresome thing is that we have to apply the substitution
 773 -- to the rules inside each Id.  Grr.  But it matters.
 774
 775 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
 776 cloneVar TopLevel env v ctxt_lvl dest_lvl
 777   = returnUs (env, v)   -- Don't clone top level things
 778 cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
 779   = ASSERT( isId v )
 780     getUs       `thenLvl` \ us ->
 781     let
 782       (subst', v1) = substAndCloneId subst us v
 783       v2           = zap_demand ctxt_lvl dest_lvl v1
 784       env'         = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
 785     in
 786     returnUs (env', v2)
 787
 788 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
 789 cloneRecVars TopLevel env vs ctxt_lvl dest_lvl
 790   = returnUs (env, vs)  -- Don't clone top level things
 791 cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
 792   = ASSERT( all isId vs )
 793     getUs                       `thenLvl` \ us ->
 794     let
 795       (subst', vs1) = substAndCloneRecIds subst us vs
 796       vs2           = map (zap_demand ctxt_lvl dest_lvl) vs1
 797       env'          = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
 798     in
 799     returnUs (env', vs2)
 800
 801         -- VERY IMPORTANT: we must zap the demand info
 802         -- if the thing is going to float out past a lambda
 803 zap_demand dest_lvl ctxt_lvl id
 804   | ctxt_lvl == dest_lvl = id                   -- Stays put
 805   | otherwise            = zapDemandIdInfo id   -- Floats out
 806 \end{code}
 807