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