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/Commentary/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
  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 Note [FloatOut inside INLINE]
 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 [] = return []
 232     do_them (b:bs) = do
 233         (lvld_bind, _) <- lvlTopBind init_env b
 234         lvld_binds <- do_them bs
 235         return (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)   = return (Type ty)
 277 lvlExpr _ env (_, AnnVar v)   = return (lookupVar env v)
 278 lvlExpr _ env (_, AnnLit lit) = return (Lit lit)
 279
 280 lvlExpr ctxt_lvl env (_, AnnApp fun arg) = do
 281     fun' <- lvl_fun fun
 282     arg' <- lvlMFE  False ctxt_lvl env arg
 283     return (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) = do
 292 -- Don't float anything out of an InlineMe; hence the iNLINE_CTXT
 293     expr' <- lvlExpr iNLINE_CTXT env expr
 294     return (Note InlineMe expr')
 295
 296 lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
 297     expr' <- lvlExpr ctxt_lvl env expr
 298     return (Note note expr')
 299
 300 lvlExpr ctxt_lvl env (_, AnnCast expr co) = do
 301     expr' <- lvlExpr ctxt_lvl env expr
 302     return (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) = do
 312     new_body <- lvlMFE True new_lvl new_env body
 313     return (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) = do
 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     rhs'  <- lvlExpr ctxt_lvl env rhs
 335     body' <- lvlExpr incd_lvl env' body
 336     return (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) = do
 343     (bind', new_env) <- lvlBind NotTopLevel ctxt_lvl env bind
 344     body' <- lvlExpr ctxt_lvl new_env body
 345     return (Let bind' body')
 346
 347 lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr ty alts) = do
 348     expr' <- lvlMFE True ctxt_lvl env expr
 349     let alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
 350     alts' <- mapM (lvl_alt alts_env) alts
 351     return (Case expr' (TB case_bndr incd_lvl) ty alts')
 352   where
 353       incd_lvl  = incMinorLvl ctxt_lvl
 354
 355       lvl_alt alts_env (con, bs, rhs) = do
 356           rhs' <- lvlMFE True incd_lvl new_env rhs
 357           return (con, bs', rhs')
 358         where
 359           bs'     = [ TB b incd_lvl | b <- bs ]
 360           new_env = extendLvlEnv alts_env bs'
 361 \end{code}
 362
 363 @lvlMFE@ is just like @lvlExpr@, except that it might let-bind
 364 the expression, so that it can itself be floated.
 365
 366 [NOTE: unlifted MFEs]
 367 We don't float unlifted MFEs, which potentially loses big opportunites.
 368 For example:
 369         \x -> f (h y)
 370 where h :: Int -> Int# is expensive. We'd like to float the (h y) outside
 371 the \x, but we don't because it's unboxed.  Possible solution: box it.
 372
 373 \begin{code}
 374 lvlMFE ::  Bool                 -- True <=> strict context [body of case or let]
 375         -> Level                -- Level of innermost enclosing lambda/tylam
 376         -> LevelEnv             -- Level of in-scope names/tyvars
 377         -> CoreExprWithFVs      -- input expression
 378         -> LvlM LevelledExpr    -- Result expression
 379
 380 lvlMFE strict_ctxt ctxt_lvl env (_, AnnType ty)
 381   = return (Type ty)
 382
 383
 384 lvlMFE strict_ctxt ctxt_lvl env ann_expr@(fvs, _)
 385   |  isUnLiftedType ty                  -- Can't let-bind it; see [NOTE: unlifted MFEs]
 386   || isInlineCtxt ctxt_lvl              -- Don't float out of an __inline__ context
 387   || exprIsTrivial expr                 -- Never float if it's trivial
 388   || not good_destination
 389   =     -- Don't float it out
 390     lvlExpr ctxt_lvl env ann_expr
 391
 392   | otherwise   -- Float it out!
 393   = do expr' <- lvlFloatRhs abs_vars dest_lvl env ann_expr
 394        var <- newLvlVar "lvl" abs_vars ty
 395        return (Let (NonRec (TB var dest_lvl) expr')
 396                    (mkVarApps (Var var) abs_vars))
 397   where
 398     expr     = deAnnotate ann_expr
 399     ty       = exprType expr
 400     dest_lvl = destLevel env fvs (isFunction ann_expr)
 401     abs_vars = abstractVars dest_lvl env fvs
 402
 403         -- A decision to float entails let-binding this thing, and we only do
 404         -- that if we'll escape a value lambda, or will go to the top level.
 405     good_destination
 406         | dest_lvl `ltMajLvl` ctxt_lvl          -- Escapes a value lambda
 407         = True
 408         -- OLD CODE: not (exprIsCheap expr) || isTopLvl dest_lvl
 409         --           see Note [Escaping a value lambda]
 410
 411         | otherwise             -- Does not escape a value lambda
 412         = isTopLvl dest_lvl     -- Only float if we are going to the top level
 413         && floatConsts env      --   and the floatConsts flag is on
 414         && not strict_ctxt      -- Don't float from a strict context
 415           -- We are keen to float something to the top level, even if it does not
 416           -- escape a lambda, because then it needs no allocation.  But it's controlled
 417           -- by a flag, because doing this too early loses opportunities for RULES
 418           -- which (needless to say) are important in some nofib programs
 419           -- (gcd is an example).
 420           --
 421           -- Beware:
 422           --    concat = /\ a -> foldr ..a.. (++) []
 423           -- was getting turned into
 424           --    concat = /\ a -> lvl a
 425           --    lvl    = /\ a -> foldr ..a.. (++) []
 426           -- which is pretty stupid.  Hence the strict_ctxt test
 427 \end{code}
 428
 429 Note [Escaping a value lambda]
 430 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 431 We want to float even cheap expressions out of value lambdas,
 432 because that saves allocation.  Consider
 433         f = \x.  .. (\y.e) ...
 434 Then we'd like to avoid allocating the (\y.e) every time we call f,
 435 (assuming e does not mention x).
 436
 437 An example where this really makes a difference is simplrun009.
 438
 439 Another reason it's good is because it makes SpecContr fire on functions.
 440 Consider
 441         f = \x. ....(f (\y.e))....
 442 After floating we get
 443         lvl = \y.e
 444         f = \x. ....(f lvl)...
 445 and that is much easier for SpecConstr to generate a robust specialisation for.
 446
 447 The OLD CODE (given where this Note is referred to) prevents floating
 448 of the example above, so I just don't understand the old code.  I
 449 don't understand the old comment either (which appears below).  I
 450 measured the effect on nofib of changing OLD CODE to 'True', and got
 451 zeros everywhere, but a 4% win for 'puzzle'.  Very small 0.5% loss for
 452 'cse'; turns out to be because our arity analysis isn't good enough
 453 yet (mentioned in Simon-nofib-notes).
 454
 455 OLD comment was:
 456          Even if it escapes a value lambda, we only
 457          float if it's not cheap (unless it'll get all the
 458          way to the top).  I've seen cases where we
 459          float dozens of tiny free expressions, which cost
 460          more to allocate than to evaluate.
 461          NB: exprIsCheap is also true of bottom expressions, which
 462              is good; we don't want to share them
 463
 464         It's only Really Bad to float a cheap expression out of a
 465         strict context, because that builds a thunk that otherwise
 466         would never be built.  So another alternative would be to
 467         add
 468                 || (strict_ctxt && not (exprIsBottom expr))
 469         to the condition above. We should really try this out.
 470
 471
 472 %************************************************************************
 473 %*                                                                      *
 474 \subsection{Bindings}
 475 %*                                                                      *
 476 %************************************************************************
 477
 478 The binding stuff works for top level too.
 479
 480 \begin{code}
 481 lvlBind :: TopLevelFlag         -- Used solely to decide whether to clone
 482         -> Level                -- Context level; might be Top even for bindings nested in the RHS
 483                                 -- of a top level binding
 484         -> LevelEnv
 485         -> CoreBindWithFVs
 486         -> LvlM (LevelledBind, LevelEnv)
 487
 488 lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
 489   | isInlineCtxt ctxt_lvl               -- Don't do anything inside InlineMe
 490   = do rhs' <- lvlExpr ctxt_lvl env rhs
 491        return (NonRec (TB bndr ctxt_lvl) rhs', env)
 492
 493   | null abs_vars
 494   = do  -- No type abstraction; clone existing binder
 495        rhs' <- lvlExpr dest_lvl env rhs
 496        (env', bndr') <- cloneVar top_lvl env bndr ctxt_lvl dest_lvl
 497        return (NonRec (TB bndr' dest_lvl) rhs', env')
 498
 499   | otherwise
 500   = do  -- Yes, type abstraction; create a new binder, extend substitution, etc
 501        rhs' <- lvlFloatRhs abs_vars dest_lvl env rhs
 502        (env', [bndr']) <- newPolyBndrs dest_lvl env abs_vars [bndr]
 503        return (NonRec (TB bndr' dest_lvl) rhs', env')
 504
 505   where
 506     bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
 507     abs_vars = abstractVars dest_lvl env bind_fvs
 508     dest_lvl = destLevel env bind_fvs (isFunction rhs)
 509 \end{code}
 510
 511
 512 \begin{code}
 513 lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
 514   | isInlineCtxt ctxt_lvl       -- Don't do anything inside InlineMe
 515   = do rhss' <- mapM (lvlExpr ctxt_lvl env) rhss
 516        return (Rec ([TB b ctxt_lvl | b <- bndrs] `zip` rhss'), env)
 517
 518   | null abs_vars
 519   = do (new_env, new_bndrs) <- cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl
 520        new_rhss <- mapM (lvlExpr ctxt_lvl new_env) rhss
 521        return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
 522
 523   | isSingleton pairs && count isId abs_vars > 1
 524   = do  -- Special case for self recursion where there are
 525         -- several variables carried around: build a local loop:
 526         --      poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
 527         -- This just makes the closures a bit smaller.  If we don't do
 528         -- this, allocation rises significantly on some programs
 529         --
 530         -- We could elaborate it for the case where there are several
 531         -- mutually functions, but it's quite a bit more complicated
 532         --
 533         -- This all seems a bit ad hoc -- sigh
 534     let
 535         (bndr,rhs) = head pairs
 536         (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 537         rhs_env = extendLvlEnv env abs_vars_w_lvls
 538     (rhs_env', new_bndr) <- cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl
 539     let
 540         (lam_bndrs, rhs_body)     = collectAnnBndrs rhs
 541         (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
 542         body_env                  = extendLvlEnv rhs_env' new_lam_bndrs
 543     new_rhs_body <- lvlExpr body_lvl body_env rhs_body
 544     (poly_env, [poly_bndr]) <- newPolyBndrs dest_lvl env abs_vars [bndr]
 545     return (Rec [(TB poly_bndr dest_lvl,
 546                mkLams abs_vars_w_lvls $
 547                mkLams new_lam_bndrs $
 548                Let (Rec [(TB new_bndr rhs_lvl, mkLams new_lam_bndrs new_rhs_body)])
 549                    (mkVarApps (Var new_bndr) lam_bndrs))],
 550                poly_env)
 551
 552   | otherwise = do  -- Non-null abs_vars
 553     (new_env, new_bndrs) <- newPolyBndrs dest_lvl env abs_vars bndrs
 554     new_rhss <- mapM (lvlFloatRhs abs_vars dest_lvl new_env) rhss
 555     return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
 556
 557   where
 558     (bndrs,rhss) = unzip pairs
 559
 560         -- Finding the free vars of the binding group is annoying
 561     bind_fvs        = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
 562                                     | (bndr, (rhs_fvs,_)) <- pairs])
 563                       `minusVarSet`
 564                       mkVarSet bndrs
 565
 566     dest_lvl = destLevel env bind_fvs (all isFunction rhss)
 567     abs_vars = abstractVars dest_lvl env bind_fvs
 568
 569 ----------------------------------------------------
 570 -- Three help functons for the type-abstraction case
 571
 572 lvlFloatRhs abs_vars dest_lvl env rhs = do
 573     rhs' <- lvlExpr rhs_lvl rhs_env rhs
 574     return (mkLams abs_vars_w_lvls rhs')
 575   where
 576     (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 577     rhs_env = extendLvlEnv env abs_vars_w_lvls
 578 \end{code}
 579
 580
 581 %************************************************************************
 582 %*                                                                      *
 583 \subsection{Deciding floatability}
 584 %*                                                                      *
 585 %************************************************************************
 586
 587 \begin{code}
 588 lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [TaggedBndr Level])
 589 -- Compute the levels for the binders of a lambda group
 590 -- The binders returned are exactly the same as the ones passed,
 591 -- but they are now paired with a level
 592 lvlLamBndrs lvl []
 593   = (lvl, [])
 594
 595 lvlLamBndrs lvl bndrs
 596   = go  (incMinorLvl lvl)
 597         False   -- Havn't bumped major level in this group
 598         [] bndrs
 599   where
 600     go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
 601         | isId bndr &&                  -- Go to the next major level if this is a value binder,
 602           not bumped_major &&           -- and we havn't already gone to the next level (one jump per group)
 603           not (isOneShotLambda bndr)    -- and it isn't a one-shot lambda
 604         = go new_lvl True (TB bndr new_lvl : rev_lvld_bndrs) bndrs
 605
 606         | otherwise
 607         = go old_lvl bumped_major (TB bndr old_lvl : rev_lvld_bndrs) bndrs
 608
 609         where
 610           new_lvl = incMajorLvl old_lvl
 611
 612     go old_lvl _ rev_lvld_bndrs []
 613         = (old_lvl, reverse rev_lvld_bndrs)
 614         -- a lambda like this (\x -> coerce t (\s -> ...))
 615         -- This happens quite a bit in state-transformer programs
 616 \end{code}
 617
 618 \begin{code}
 619   -- Destintion level is the max Id level of the expression
 620   -- (We'll abstract the type variables, if any.)
 621 destLevel :: LevelEnv -> VarSet -> Bool -> Level
 622 destLevel env fvs is_function
 623   |  floatLams env
 624   && is_function = tOP_LEVEL            -- Send functions to top level; see
 625                                         -- the comments with isFunction
 626   | otherwise    = maxIdLevel env fvs
 627
 628 isFunction :: CoreExprWithFVs -> Bool
 629 -- The idea here is that we want to float *functions* to
 630 -- the top level.  This saves no work, but
 631 --      (a) it can make the host function body a lot smaller,
 632 --              and hence inlinable.
 633 --      (b) it can also save allocation when the function is recursive:
 634 --          h = \x -> letrec f = \y -> ...f...y...x...
 635 --                    in f x
 636 --     becomes
 637 --          f = \x y -> ...(f x)...y...x...
 638 --          h = \x -> f x x
 639 --     No allocation for f now.
 640 -- We may only want to do this if there are sufficiently few free
 641 -- variables.  We certainly only want to do it for values, and not for
 642 -- constructors.  So the simple thing is just to look for lambdas
 643 isFunction (_, AnnLam b e) | isId b    = True
 644                            | otherwise = isFunction e
 645 isFunction (_, AnnNote n e)            = isFunction e
 646 isFunction other                       = False
 647 \end{code}
 648
 649
 650 %************************************************************************
 651 %*                                                                      *
 652 \subsection{Free-To-Level Monad}
 653 %*                                                                      *
 654 %************************************************************************
 655
 656 \begin{code}
 657 type LevelEnv = (FloatOutSwitches,
 658                  VarEnv Level,                  -- Domain is *post-cloned* TyVars and Ids
 659                  Subst,                         -- Domain is pre-cloned Ids; tracks the in-scope set
 660                                                 --      so that subtitution is capture-avoiding
 661                  IdEnv ([Var], LevelledExpr))   -- Domain is pre-cloned Ids
 662         -- We clone let-bound variables so that they are still
 663         -- distinct when floated out; hence the SubstEnv/IdEnv.
 664         -- (see point 3 of the module overview comment).
 665         -- We also use these envs when making a variable polymorphic
 666         -- because we want to float it out past a big lambda.
 667         --
 668         -- The Subst and IdEnv always implement the same mapping, but the
 669         -- Subst maps to CoreExpr and the IdEnv to LevelledExpr
 670         -- Since the range is always a variable or type application,
 671         -- there is never any difference between the two, but sadly
 672         -- the types differ.  The SubstEnv is used when substituting in
 673         -- a variable's IdInfo; the IdEnv when we find a Var.
 674         --
 675         -- In addition the IdEnv records a list of tyvars free in the
 676         -- type application, just so we don't have to call freeVars on
 677         -- the type application repeatedly.
 678         --
 679         -- The domain of the both envs is *pre-cloned* Ids, though
 680         --
 681         -- The domain of the VarEnv Level is the *post-cloned* Ids
 682
 683 initialEnv :: FloatOutSwitches -> LevelEnv
 684 initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
 685
 686 floatLams :: LevelEnv -> Bool
 687 floatLams (FloatOutSw float_lams _, _, _, _) = float_lams
 688
 689 floatConsts :: LevelEnv -> Bool
 690 floatConsts (FloatOutSw _ float_consts, _, _, _) = float_consts
 691
 692 extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
 693 -- Used when *not* cloning
 694 extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
 695   = (float_lams,
 696      foldl add_lvl lvl_env prs,
 697      foldl del_subst subst prs,
 698      foldl del_id id_env prs)
 699   where
 700     add_lvl   env (TB v l) = extendVarEnv env v l
 701     del_subst env (TB v _) = extendInScope env v
 702     del_id    env (TB v _) = delVarEnv env v
 703   -- We must remove any clone for this variable name in case of
 704   -- shadowing.  This bit me in the following case
 705   -- (in nofib/real/gg/Spark.hs):
 706   --
 707   --   case ds of wild {
 708   --     ... -> case e of wild {
 709   --              ... -> ... wild ...
 710   --            }
 711   --   }
 712   --
 713   -- The inside occurrence of @wild@ was being replaced with @ds@,
 714   -- incorrectly, because the SubstEnv was still lying around.  Ouch!
 715   -- KSW 2000-07.
 716
 717 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
 718 -- (see point 4 of the module overview comment)
 719 extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
 720   = (float_lams,
 721      extendVarEnv lvl_env case_bndr lvl,
 722      extendIdSubst subst case_bndr (Var scrut_var),
 723      extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
 724
 725 extendCaseBndrLvlEnv env scrut case_bndr lvl
 726   = extendLvlEnv          env [TB case_bndr lvl]
 727
 728 extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
 729   = (float_lams,
 730      foldl add_lvl   lvl_env bndr_pairs,
 731      foldl add_subst subst   bndr_pairs,
 732      foldl add_id    id_env  bndr_pairs)
 733   where
 734      add_lvl   env (v,v') = extendVarEnv env v' dest_lvl
 735      add_subst env (v,v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
 736      add_id    env (v,v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
 737
 738 extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
 739   = (float_lams,
 740      foldl add_lvl   lvl_env bndr_pairs,
 741      new_subst,
 742      foldl add_id    id_env  bndr_pairs)
 743   where
 744      add_lvl   env (v,v') = extendVarEnv env v' lvl
 745      add_id    env (v,v') = extendVarEnv env v ([v'], Var v')
 746
 747
 748 maxIdLevel :: LevelEnv -> VarSet -> Level
 749 maxIdLevel (_, lvl_env,_,id_env) var_set
 750   = foldVarSet max_in tOP_LEVEL var_set
 751   where
 752     max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
 753                                                 Just (abs_vars, _) -> abs_vars
 754                                                 Nothing            -> [in_var])
 755
 756     max_out out_var lvl
 757         | isId out_var = case lookupVarEnv lvl_env out_var of
 758                                 Just lvl' -> maxLvl lvl' lvl
 759                                 Nothing   -> lvl
 760         | otherwise    = lvl    -- Ignore tyvars in *maxIdLevel*
 761
 762 lookupVar :: LevelEnv -> Id -> LevelledExpr
 763 lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
 764                                        Just (_, expr) -> expr
 765                                        other          -> Var v
 766
 767 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
 768         -- Find the variables in fvs, free vars of the target expresion,
 769         -- whose level is greater than the destination level
 770         -- These are the ones we are going to abstract out
 771 abstractVars dest_lvl (_, lvl_env, _, id_env) fvs
 772   = map zap $ uniq $ sortLe le
 773         [var | fv <- varSetElems fvs
 774              , var <- absVarsOf id_env fv
 775              , abstract_me var ]
 776         -- NB: it's important to call abstract_me only on the OutIds the
 777         -- come from absVarsOf (not on fv, which is an InId)
 778   where
 779         -- Sort the variables so the true type variables come first;
 780         -- the tyvars scope over Ids and coercion vars
 781     v1 `le` v2 = case (is_tv v1, is_tv v2) of
 782                    (True, False) -> True
 783                    (False, True) -> False
 784                    other         -> v1 <= v2    -- Same family
 785
 786     is_tv v = isTyVar v && not (isCoVar v)
 787
 788     uniq :: [Var] -> [Var]
 789         -- Remove adjacent duplicates; the sort will have brought them together
 790     uniq (v1:v2:vs) | v1 == v2  = uniq (v2:vs)
 791                     | otherwise = v1 : uniq (v2:vs)
 792     uniq vs = vs
 793
 794     abstract_me v = case lookupVarEnv lvl_env v of
 795                         Just lvl -> dest_lvl `ltLvl` lvl
 796                         Nothing  -> False
 797
 798         -- We are going to lambda-abstract, so nuke any IdInfo,
 799         -- and add the tyvars of the Id (if necessary)
 800     zap v | isId v = WARN( workerExists (idWorkerInfo v) ||
 801                            not (isEmptySpecInfo (idSpecialisation v)),
 802                            text "absVarsOf: discarding info on" <+> ppr v )
 803                      setIdInfo v vanillaIdInfo
 804           | otherwise = v
 805
 806 absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> [Var]
 807         -- If f is free in the expression, and f maps to poly_f a b c in the
 808         -- current substitution, then we must report a b c as candidate type
 809         -- variables
 810         --
 811         -- Also, if x::a is an abstracted variable, then so is a; that is,
 812         --      we must look in x's type
 813         -- And similarly if x is a coercion variable.
 814 absVarsOf id_env v
 815   | isId v    = [av2 | av1 <- lookup_avs v
 816                      , av2 <- add_tyvars av1]
 817   | isCoVar v = add_tyvars v
 818   | otherwise = [v]
 819
 820   where
 821     lookup_avs v = case lookupVarEnv id_env v of
 822                         Just (abs_vars, _) -> abs_vars
 823                         Nothing            -> [v]
 824
 825     add_tyvars v = v : varSetElems (varTypeTyVars v)
 826 \end{code}
 827
 828 \begin{code}
 829 type LvlM result = UniqSM result
 830
 831 initLvl         = initUs_
 832 \end{code}
 833
 834 \begin{code}
 835 newPolyBndrs dest_lvl env abs_vars bndrs = do
 836     uniqs <- getUniquesM
 837     let new_bndrs = zipWith mk_poly_bndr bndrs uniqs
 838     return (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
 839   where
 840     mk_poly_bndr bndr uniq = mkSysLocal (mkFastString str) uniq poly_ty
 841                            where
 842                              str     = "poly_" ++ occNameString (getOccName bndr)
 843                              poly_ty = mkPiTypes abs_vars (idType bndr)
 844
 845
 846 newLvlVar :: String
 847           -> [CoreBndr] -> Type         -- Abstract wrt these bndrs
 848           -> LvlM Id
 849 newLvlVar str vars body_ty = do
 850     uniq <- getUniqueM
 851     return (mkSysLocal (mkFastString str) uniq (mkPiTypes vars body_ty))
 852
 853 -- The deeply tiresome thing is that we have to apply the substitution
 854 -- to the rules inside each Id.  Grr.  But it matters.
 855
 856 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
 857 cloneVar TopLevel env v ctxt_lvl dest_lvl
 858   = return (env, v)     -- Don't clone top level things
 859 cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
 860   = ASSERT( isId v ) do
 861     us <- getUniqueSupplyM
 862     let
 863       (subst', v1) = cloneIdBndr subst us v
 864       v2           = zap_demand ctxt_lvl dest_lvl v1
 865       env'         = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
 866     return (env', v2)
 867
 868 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
 869 cloneRecVars TopLevel env vs ctxt_lvl dest_lvl
 870   = return (env, vs)    -- Don't clone top level things
 871 cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
 872   = ASSERT( all isId vs ) do
 873     us <- getUniqueSupplyM
 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     return (env', vs2)
 879
 880         -- VERY IMPORTANT: we must zap the demand info
 881         -- if the thing is going to float out past a lambda,
 882         -- or if it's going to top level (where things can't be strict)
 883 zap_demand dest_lvl ctxt_lvl id
 884   | ctxt_lvl == dest_lvl,
 885     not (isTopLvl dest_lvl) = id        -- Stays, and not going to top level
 886   | otherwise               = zapDemandIdInfo id        -- Floats out
 887 \end{code}
 888