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