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