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, transferPolyIdInfo,
  72                           idSpecialisation, idWorkerInfo, setIdInfo
  73                         )
  74 import IdInfo           ( workerExists, vanillaIdInfo, isEmptySpecInfo,
  75                           setNewStrictnessInfo, newStrictnessInfo,
  76                           setArityInfo, arityInfo )
  77 import Var
  78 import VarSet
  79 import VarEnv
  80 import Name             ( getOccName )
  81 import OccName          ( occNameString )
  82 import Type             ( isUnLiftedType, Type )
  83 import BasicTypes       ( TopLevelFlag(..) )
  84 import UniqSupply
  85 import Util             ( sortLe, isSingleton, count )
  86 import Outputable
  87 import FastString
  88 \end{code}
  89
  90 %************************************************************************
  91 %*                                                                      *
  92 \subsection{Level numbers}
  93 %*                                                                      *
  94 %************************************************************************
  95
  96 \begin{code}
  97 data Level = InlineCtxt -- A level that's used only for
  98                         -- the context parameter ctxt_lvl
  99            | Level Int  -- Level number of enclosing lambdas
 100                    Int  -- Number of big-lambda and/or case expressions between
 101                         -- here and the nearest enclosing lambda
 102 \end{code}
 103
 104 The {\em level number} on a (type-)lambda-bound variable is the
 105 nesting depth of the (type-)lambda which binds it.  The outermost lambda
 106 has level 1, so (Level 0 0) means that the variable is bound outside any lambda.
 107
 108 On an expression, it's the maximum level number of its free
 109 (type-)variables.  On a let(rec)-bound variable, it's the level of its
 110 RHS.  On a case-bound variable, it's the number of enclosing lambdas.
 111
 112 Top-level variables: level~0.  Those bound on the RHS of a top-level
 113 definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
 114 as ``subscripts'')...
 115 \begin{verbatim}
 116 a_0 = let  b_? = ...  in
 117            x_1 = ... b ... in ...
 118 \end{verbatim}
 119
 120 The main function @lvlExpr@ carries a ``context level'' (@ctxt_lvl@).
 121 That's meant to be the level number of the enclosing binder in the
 122 final (floated) program.  If the level number of a sub-expression is
 123 less than that of the context, then it might be worth let-binding the
 124 sub-expression so that it will indeed float.
 125
 126 If you can float to level @Level 0 0@ worth doing so because then your
 127 allocation becomes static instead of dynamic.  We always start with
 128 context @Level 0 0@.
 129
 130
 131 Note [FloatOut inside INLINE]
 132 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 133 @InlineCtxt@ very similar to @Level 0 0@, but is used for one purpose:
 134 to say "don't float anything out of here".  That's exactly what we
 135 want for the body of an INLINE, where we don't want to float anything
 136 out at all.  See notes with lvlMFE below.
 137
 138 But, check this out:
 139
 140 -- At one time I tried the effect of not float anything out of an InlineMe,
 141 -- but it sometimes works badly.  For example, consider PrelArr.done.  It
 142 -- has the form         __inline (\d. e)
 143 -- where e doesn't mention d.  If we float this to
 144 --      __inline (let x = e in \d. x)
 145 -- things are bad.  The inliner doesn't even inline it because it doesn't look
 146 -- like a head-normal form.  So it seems a lesser evil to let things float.
 147 -- In SetLevels we do set the context to (Level 0 0) when we get to an InlineMe
 148 -- which discourages floating out.
 149
 150 So the conclusion is: don't do any floating at all inside an InlineMe.
 151 (In the above example, don't float the {x=e} out of the \d.)
 152
 153 One particular case is that of workers: we don't want to float the
 154 call to the worker outside the wrapper, otherwise the worker might get
 155 inlined into the floated expression, and an importing module won't see
 156 the worker at all.
 157
 158 \begin{code}
 159 type LevelledExpr  = TaggedExpr Level
 160 type LevelledBind  = TaggedBind Level
 161
 162 tOP_LEVEL   = Level 0 0
 163 iNLINE_CTXT = InlineCtxt
 164
 165 incMajorLvl :: Level -> Level
 166 -- For InlineCtxt we ignore any inc's; we don't want
 167 -- to do any floating at all; see notes above
 168 incMajorLvl InlineCtxt          = InlineCtxt
 169 incMajorLvl (Level major minor) = Level (major+1) 0
 170
 171 incMinorLvl :: Level -> Level
 172 incMinorLvl InlineCtxt          = InlineCtxt
 173 incMinorLvl (Level major minor) = Level major (minor+1)
 174
 175 maxLvl :: Level -> Level -> Level
 176 maxLvl InlineCtxt l2  = l2
 177 maxLvl l1  InlineCtxt = l1
 178 maxLvl l1@(Level maj1 min1) l2@(Level maj2 min2)
 179   | (maj1 > maj2) || (maj1 == maj2 && min1 > min2) = l1
 180   | otherwise                                      = l2
 181
 182 ltLvl :: Level -> Level -> Bool
 183 ltLvl any_lvl    InlineCtxt  = False
 184 ltLvl InlineCtxt (Level _ _) = True
 185 ltLvl (Level maj1 min1) (Level maj2 min2)
 186   = (maj1 < maj2) || (maj1 == maj2 && min1 < min2)
 187
 188 ltMajLvl :: Level -> Level -> Bool
 189     -- Tells if one level belongs to a difft *lambda* level to another
 190 ltMajLvl any_lvl        InlineCtxt     = False
 191 ltMajLvl InlineCtxt     (Level maj2 _) = 0 < maj2
 192 ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2
 193
 194 isTopLvl :: Level -> Bool
 195 isTopLvl (Level 0 0) = True
 196 isTopLvl other       = False
 197
 198 isInlineCtxt :: Level -> Bool
 199 isInlineCtxt InlineCtxt = True
 200 isInlineCtxt other      = False
 201
 202 instance Outputable Level where
 203   ppr InlineCtxt      = text "<INLINE>"
 204   ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]
 205
 206 instance Eq Level where
 207   InlineCtxt        == InlineCtxt        = True
 208   (Level maj1 min1) == (Level maj2 min2) = maj1==maj2 && min1==min2
 209   l1                == l2                = False
 210 \end{code}
 211
 212
 213 %************************************************************************
 214 %*                                                                      *
 215 \subsection{Main level-setting code}
 216 %*                                                                      *
 217 %************************************************************************
 218
 219 \begin{code}
 220 setLevels :: FloatOutSwitches
 221           -> [CoreBind]
 222           -> UniqSupply
 223           -> [LevelledBind]
 224
 225 setLevels float_lams binds us
 226   = initLvl us (do_them binds)
 227   where
 228     -- "do_them"'s main business is to thread the monad along
 229     -- It gives each top binding the same empty envt, because
 230     -- things unbound in the envt have level number zero implicitly
 231     do_them :: [CoreBind] -> LvlM [LevelledBind]
 232
 233     do_them [] = return []
 234     do_them (b:bs) = do
 235         (lvld_bind, _) <- lvlTopBind init_env b
 236         lvld_binds <- do_them bs
 237         return (lvld_bind : lvld_binds)
 238
 239     init_env = initialEnv float_lams
 240
 241 lvlTopBind env (NonRec binder rhs)
 242   = lvlBind TopLevel tOP_LEVEL env (AnnNonRec binder (freeVars rhs))
 243                                         -- Rhs can have no free vars!
 244
 245 lvlTopBind env (Rec pairs)
 246   = lvlBind TopLevel tOP_LEVEL env (AnnRec [(b,freeVars rhs) | (b,rhs) <- pairs])
 247 \end{code}
 248
 249 %************************************************************************
 250 %*                                                                      *
 251 \subsection{Setting expression levels}
 252 %*                                                                      *
 253 %************************************************************************
 254
 255 \begin{code}
 256 lvlExpr :: Level                -- ctxt_lvl: Level of enclosing expression
 257         -> LevelEnv             -- Level of in-scope names/tyvars
 258         -> CoreExprWithFVs      -- input expression
 259         -> LvlM LevelledExpr    -- Result expression
 260 \end{code}
 261
 262 The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
 263 binder.  Here's an example
 264
 265         v = \x -> ...\y -> let r = case (..x..) of
 266                                         ..x..
 267                            in ..
 268
 269 When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
 270 the level of @r@, even though it's inside a level-2 @\y@.  It's
 271 important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
 272 don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
 273 --- because it isn't a *maximal* free expression.
 274
 275 If there were another lambda in @r@'s rhs, it would get level-2 as well.
 276
 277 \begin{code}
 278 lvlExpr _ _ (_, AnnType ty)   = return (Type ty)
 279 lvlExpr _ env (_, AnnVar v)   = return (lookupVar env v)
 280 lvlExpr _ env (_, AnnLit lit) = return (Lit lit)
 281
 282 lvlExpr ctxt_lvl env (_, AnnApp fun arg) = do
 283     fun' <- lvl_fun fun
 284     arg' <- lvlMFE  False ctxt_lvl env arg
 285     return (App fun' arg')
 286   where
 287 -- gaw 2004
 288     lvl_fun (_, AnnCase _ _ _ _) = lvlMFE True ctxt_lvl env fun
 289     lvl_fun other              = lvlExpr ctxt_lvl env fun
 290         -- We don't do MFE on partial applications generally,
 291         -- but we do if the function is big and hairy, like a case
 292
 293 lvlExpr ctxt_lvl env (_, AnnNote InlineMe expr) = do
 294 -- Don't float anything out of an InlineMe; hence the iNLINE_CTXT
 295     expr' <- lvlExpr iNLINE_CTXT env expr
 296     return (Note InlineMe expr')
 297
 298 lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
 299     expr' <- lvlExpr ctxt_lvl env expr
 300     return (Note note expr')
 301
 302 lvlExpr ctxt_lvl env (_, AnnCast expr co) = do
 303     expr' <- lvlExpr ctxt_lvl env expr
 304     return (Cast expr' co)
 305
 306 -- We don't split adjacent lambdas.  That is, given
 307 --      \x y -> (x+1,y)
 308 -- we don't float to give
 309 --      \x -> let v = x+y in \y -> (v,y)
 310 -- Why not?  Because partial applications are fairly rare, and splitting
 311 -- lambdas makes them more expensive.
 312
 313 lvlExpr ctxt_lvl env expr@(_, AnnLam bndr rhs) = do
 314     new_body <- lvlMFE True new_lvl new_env body
 315     return (mkLams new_bndrs new_body)
 316   where
 317     (bndrs, body)        = collectAnnBndrs expr
 318     (new_lvl, new_bndrs) = lvlLamBndrs ctxt_lvl bndrs
 319     new_env              = extendLvlEnv env new_bndrs
 320         -- At one time we called a special verion of collectBinders,
 321         -- which ignored coercions, because we don't want to split
 322         -- a lambda like this (\x -> coerce t (\s -> ...))
 323         -- This used to happen quite a bit in state-transformer programs,
 324         -- but not nearly so much now non-recursive newtypes are transparent.
 325         -- [See SetLevels rev 1.50 for a version with this approach.]
 326
 327 lvlExpr ctxt_lvl env (_, AnnLet (AnnNonRec bndr rhs) body)
 328   | isUnLiftedType (idType bndr) = do
 329         -- Treat unlifted let-bindings (let x = b in e) just like (case b of x -> e)
 330         -- That is, leave it exactly where it is
 331         -- We used to float unlifted bindings too (e.g. to get a cheap primop
 332         -- outside a lambda (to see how, look at lvlBind in rev 1.58)
 333         -- but an unrelated change meant that these unlifed bindings
 334         -- could get to the top level which is bad.  And there's not much point;
 335         -- unlifted bindings are always cheap, and so hardly worth floating.
 336     rhs'  <- lvlExpr ctxt_lvl env rhs
 337     body' <- lvlExpr incd_lvl env' body
 338     return (Let (NonRec bndr' rhs') body')
 339   where
 340     incd_lvl = incMinorLvl ctxt_lvl
 341     bndr' = TB bndr incd_lvl
 342     env'  = extendLvlEnv env [bndr']
 343
 344 lvlExpr ctxt_lvl env (_, AnnLet bind body) = do
 345     (bind', new_env) <- lvlBind NotTopLevel ctxt_lvl env bind
 346     body' <- lvlExpr ctxt_lvl new_env body
 347     return (Let bind' body')
 348
 349 lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr ty alts) = do
 350     expr' <- lvlMFE True ctxt_lvl env expr
 351     let alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
 352     alts' <- mapM (lvl_alt alts_env) alts
 353     return (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) = do
 358           rhs' <- lvlMFE True incd_lvl new_env rhs
 359           return (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   = return (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   = do expr' <- lvlFloatRhs abs_vars dest_lvl env ann_expr
 396        var <- newLvlVar "lvl" abs_vars ty
 397        return (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   = do rhs' <- lvlExpr ctxt_lvl env rhs
 493        return (NonRec (TB bndr ctxt_lvl) rhs', env)
 494
 495   | null abs_vars
 496   = do  -- No type abstraction; clone existing binder
 497        rhs' <- lvlExpr dest_lvl env rhs
 498        (env', bndr') <- cloneVar top_lvl env bndr ctxt_lvl dest_lvl
 499        return (NonRec (TB bndr' dest_lvl) rhs', env')
 500
 501   | otherwise
 502   = do  -- Yes, type abstraction; create a new binder, extend substitution, etc
 503        rhs' <- lvlFloatRhs abs_vars dest_lvl env rhs
 504        (env', [bndr']) <- newPolyBndrs dest_lvl env abs_vars [bndr]
 505        return (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   = do rhss' <- mapM (lvlExpr ctxt_lvl env) rhss
 518        return (Rec ([TB b ctxt_lvl | b <- bndrs] `zip` rhss'), env)
 519
 520   | null abs_vars
 521   = do (new_env, new_bndrs) <- cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl
 522        new_rhss <- mapM (lvlExpr ctxt_lvl new_env) rhss
 523        return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
 524
 525   | isSingleton pairs && count isId abs_vars > 1
 526   = do  -- 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     (rhs_env', new_bndr) <- cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl
 541     let
 542         (lam_bndrs, rhs_body)     = collectAnnBndrs rhs
 543         (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
 544         body_env                  = extendLvlEnv rhs_env' new_lam_bndrs
 545     new_rhs_body <- lvlExpr body_lvl body_env rhs_body
 546     (poly_env, [poly_bndr]) <- newPolyBndrs dest_lvl env abs_vars [bndr]
 547     return (Rec [(TB poly_bndr dest_lvl,
 548                mkLams abs_vars_w_lvls $
 549                mkLams new_lam_bndrs $
 550                Let (Rec [(TB new_bndr rhs_lvl, mkLams new_lam_bndrs new_rhs_body)])
 551                    (mkVarApps (Var new_bndr) lam_bndrs))],
 552                poly_env)
 553
 554   | otherwise = do  -- Non-null abs_vars
 555     (new_env, new_bndrs) <- newPolyBndrs dest_lvl env abs_vars bndrs
 556     new_rhss <- mapM (lvlFloatRhs abs_vars dest_lvl new_env) rhss
 557     return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
 558
 559   where
 560     (bndrs,rhss) = unzip pairs
 561
 562         -- Finding the free vars of the binding group is annoying
 563     bind_fvs        = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
 564                                     | (bndr, (rhs_fvs,_)) <- pairs])
 565                       `minusVarSet`
 566                       mkVarSet bndrs
 567
 568     dest_lvl = destLevel env bind_fvs (all isFunction rhss)
 569     abs_vars = abstractVars dest_lvl env bind_fvs
 570
 571 ----------------------------------------------------
 572 -- Three help functons for the type-abstraction case
 573
 574 lvlFloatRhs abs_vars dest_lvl env rhs = do
 575     rhs' <- lvlExpr rhs_lvl rhs_env rhs
 576     return (mkLams abs_vars_w_lvls rhs')
 577   where
 578     (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 579     rhs_env = extendLvlEnv env abs_vars_w_lvls
 580 \end{code}
 581
 582
 583 %************************************************************************
 584 %*                                                                      *
 585 \subsection{Deciding floatability}
 586 %*                                                                      *
 587 %************************************************************************
 588
 589 \begin{code}
 590 lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [TaggedBndr Level])
 591 -- Compute the levels for the binders of a lambda group
 592 -- The binders returned are exactly the same as the ones passed,
 593 -- but they are now paired with a level
 594 lvlLamBndrs lvl []
 595   = (lvl, [])
 596
 597 lvlLamBndrs lvl bndrs
 598   = go  (incMinorLvl lvl)
 599         False   -- Havn't bumped major level in this group
 600         [] bndrs
 601   where
 602     go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
 603         | isId bndr &&                  -- Go to the next major level if this is a value binder,
 604           not bumped_major &&           -- and we havn't already gone to the next level (one jump per group)
 605           not (isOneShotLambda bndr)    -- and it isn't a one-shot lambda
 606         = go new_lvl True (TB bndr new_lvl : rev_lvld_bndrs) bndrs
 607
 608         | otherwise
 609         = go old_lvl bumped_major (TB bndr old_lvl : rev_lvld_bndrs) bndrs
 610
 611         where
 612           new_lvl = incMajorLvl old_lvl
 613
 614     go old_lvl _ rev_lvld_bndrs []
 615         = (old_lvl, reverse rev_lvld_bndrs)
 616         -- a lambda like this (\x -> coerce t (\s -> ...))
 617         -- This happens quite a bit in state-transformer programs
 618 \end{code}
 619
 620 \begin{code}
 621   -- Destintion level is the max Id level of the expression
 622   -- (We'll abstract the type variables, if any.)
 623 destLevel :: LevelEnv -> VarSet -> Bool -> Level
 624 destLevel env fvs is_function
 625   |  floatLams env
 626   && is_function = tOP_LEVEL            -- Send functions to top level; see
 627                                         -- the comments with isFunction
 628   | otherwise    = maxIdLevel env fvs
 629
 630 isFunction :: CoreExprWithFVs -> Bool
 631 -- The idea here is that we want to float *functions* to
 632 -- the top level.  This saves no work, but
 633 --      (a) it can make the host function body a lot smaller,
 634 --              and hence inlinable.
 635 --      (b) it can also save allocation when the function is recursive:
 636 --          h = \x -> letrec f = \y -> ...f...y...x...
 637 --                    in f x
 638 --     becomes
 639 --          f = \x y -> ...(f x)...y...x...
 640 --          h = \x -> f x x
 641 --     No allocation for f now.
 642 -- We may only want to do this if there are sufficiently few free
 643 -- variables.  We certainly only want to do it for values, and not for
 644 -- constructors.  So the simple thing is just to look for lambdas
 645 isFunction (_, AnnLam b e) | isId b    = True
 646                            | otherwise = isFunction e
 647 isFunction (_, AnnNote n e)            = isFunction e
 648 isFunction other                       = False
 649 \end{code}
 650
 651
 652 %************************************************************************
 653 %*                                                                      *
 654 \subsection{Free-To-Level Monad}
 655 %*                                                                      *
 656 %************************************************************************
 657
 658 \begin{code}
 659 type LevelEnv = (FloatOutSwitches,
 660                  VarEnv Level,                  -- Domain is *post-cloned* TyVars and Ids
 661                  Subst,                         -- Domain is pre-cloned Ids; tracks the in-scope set
 662                                                 --      so that subtitution is capture-avoiding
 663                  IdEnv ([Var], LevelledExpr))   -- Domain is pre-cloned Ids
 664         -- We clone let-bound variables so that they are still
 665         -- distinct when floated out; hence the SubstEnv/IdEnv.
 666         -- (see point 3 of the module overview comment).
 667         -- We also use these envs when making a variable polymorphic
 668         -- because we want to float it out past a big lambda.
 669         --
 670         -- The Subst and IdEnv always implement the same mapping, but the
 671         -- Subst maps to CoreExpr and the IdEnv to LevelledExpr
 672         -- Since the range is always a variable or type application,
 673         -- there is never any difference between the two, but sadly
 674         -- the types differ.  The SubstEnv is used when substituting in
 675         -- a variable's IdInfo; the IdEnv when we find a Var.
 676         --
 677         -- In addition the IdEnv records a list of tyvars free in the
 678         -- type application, just so we don't have to call freeVars on
 679         -- the type application repeatedly.
 680         --
 681         -- The domain of the both envs is *pre-cloned* Ids, though
 682         --
 683         -- The domain of the VarEnv Level is the *post-cloned* Ids
 684
 685 initialEnv :: FloatOutSwitches -> LevelEnv
 686 initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
 687
 688 floatLams :: LevelEnv -> Bool
 689 floatLams (FloatOutSw float_lams _, _, _, _) = float_lams
 690
 691 floatConsts :: LevelEnv -> Bool
 692 floatConsts (FloatOutSw _ float_consts, _, _, _) = float_consts
 693
 694 extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
 695 -- Used when *not* cloning
 696 extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
 697   = (float_lams,
 698      foldl add_lvl lvl_env prs,
 699      foldl del_subst subst prs,
 700      foldl del_id id_env prs)
 701   where
 702     add_lvl   env (TB v l) = extendVarEnv env v l
 703     del_subst env (TB v _) = extendInScope env v
 704     del_id    env (TB v _) = delVarEnv env v
 705   -- We must remove any clone for this variable name in case of
 706   -- shadowing.  This bit me in the following case
 707   -- (in nofib/real/gg/Spark.hs):
 708   --
 709   --   case ds of wild {
 710   --     ... -> case e of wild {
 711   --              ... -> ... wild ...
 712   --            }
 713   --   }
 714   --
 715   -- The inside occurrence of @wild@ was being replaced with @ds@,
 716   -- incorrectly, because the SubstEnv was still lying around.  Ouch!
 717   -- KSW 2000-07.
 718
 719 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
 720 -- (see point 4 of the module overview comment)
 721 extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
 722   = (float_lams,
 723      extendVarEnv lvl_env case_bndr lvl,
 724      extendIdSubst subst case_bndr (Var scrut_var),
 725      extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
 726
 727 extendCaseBndrLvlEnv env scrut case_bndr lvl
 728   = extendLvlEnv          env [TB case_bndr lvl]
 729
 730 extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
 731   = (float_lams,
 732      foldl add_lvl   lvl_env bndr_pairs,
 733      foldl add_subst subst   bndr_pairs,
 734      foldl add_id    id_env  bndr_pairs)
 735   where
 736      add_lvl   env (v,v') = extendVarEnv env v' dest_lvl
 737      add_subst env (v,v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
 738      add_id    env (v,v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
 739
 740 extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
 741   = (float_lams,
 742      foldl add_lvl   lvl_env bndr_pairs,
 743      new_subst,
 744      foldl add_id    id_env  bndr_pairs)
 745   where
 746      add_lvl   env (v,v') = extendVarEnv env v' lvl
 747      add_id    env (v,v') = extendVarEnv env v ([v'], Var v')
 748
 749
 750 maxIdLevel :: LevelEnv -> VarSet -> Level
 751 maxIdLevel (_, lvl_env,_,id_env) var_set
 752   = foldVarSet max_in tOP_LEVEL var_set
 753   where
 754     max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
 755                                                 Just (abs_vars, _) -> abs_vars
 756                                                 Nothing            -> [in_var])
 757
 758     max_out out_var lvl
 759         | isId out_var = case lookupVarEnv lvl_env out_var of
 760                                 Just lvl' -> maxLvl lvl' lvl
 761                                 Nothing   -> lvl
 762         | otherwise    = lvl    -- Ignore tyvars in *maxIdLevel*
 763
 764 lookupVar :: LevelEnv -> Id -> LevelledExpr
 765 lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
 766                                        Just (_, expr) -> expr
 767                                        other          -> Var v
 768
 769 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
 770         -- Find the variables in fvs, free vars of the target expresion,
 771         -- whose level is greater than the destination level
 772         -- These are the ones we are going to abstract out
 773 abstractVars dest_lvl (_, lvl_env, _, id_env) fvs
 774   = map zap $ uniq $ sortLe le
 775         [var | fv <- varSetElems fvs
 776              , var <- absVarsOf id_env fv
 777              , abstract_me var ]
 778         -- NB: it's important to call abstract_me only on the OutIds the
 779         -- come from absVarsOf (not on fv, which is an InId)
 780   where
 781         -- Sort the variables so the true type variables come first;
 782         -- the tyvars scope over Ids and coercion vars
 783     v1 `le` v2 = case (is_tv v1, is_tv v2) of
 784                    (True, False) -> True
 785                    (False, True) -> False
 786                    other         -> v1 <= v2    -- Same family
 787
 788     is_tv v = isTyVar v && not (isCoVar v)
 789
 790     uniq :: [Var] -> [Var]
 791         -- Remove adjacent duplicates; the sort will have brought them together
 792     uniq (v1:v2:vs) | v1 == v2  = uniq (v2:vs)
 793                     | otherwise = v1 : uniq (v2:vs)
 794     uniq vs = vs
 795
 796     abstract_me v = case lookupVarEnv lvl_env v of
 797                         Just lvl -> dest_lvl `ltLvl` lvl
 798                         Nothing  -> False
 799
 800         -- We are going to lambda-abstract, so nuke any IdInfo,
 801         -- and add the tyvars of the Id (if necessary)
 802     zap v | isId v = WARN( workerExists (idWorkerInfo v) ||
 803                            not (isEmptySpecInfo (idSpecialisation v)),
 804                            text "absVarsOf: discarding info on" <+> ppr v )
 805                      setIdInfo v vanillaIdInfo
 806           | otherwise = v
 807
 808 absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> [Var]
 809         -- If f is free in the expression, and f maps to poly_f a b c in the
 810         -- current substitution, then we must report a b c as candidate type
 811         -- variables
 812         --
 813         -- Also, if x::a is an abstracted variable, then so is a; that is,
 814         --      we must look in x's type
 815         -- And similarly if x is a coercion variable.
 816 absVarsOf id_env v
 817   | isId v    = [av2 | av1 <- lookup_avs v
 818                      , av2 <- add_tyvars av1]
 819   | isCoVar v = add_tyvars v
 820   | otherwise = [v]
 821
 822   where
 823     lookup_avs v = case lookupVarEnv id_env v of
 824                         Just (abs_vars, _) -> abs_vars
 825                         Nothing            -> [v]
 826
 827     add_tyvars v = v : varSetElems (varTypeTyVars v)
 828 \end{code}
 829
 830 \begin{code}
 831 type LvlM result = UniqSM result
 832
 833 initLvl         = initUs_
 834 \end{code}
 835
 836
 837 \begin{code}
 838 newPolyBndrs dest_lvl env abs_vars bndrs = do
 839     uniqs <- getUniquesM
 840     let new_bndrs = zipWith mk_poly_bndr bndrs uniqs
 841     return (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
 842   where
 843     mk_poly_bndr bndr uniq = transferPolyIdInfo bndr $  -- Note [transferPolyIdInfo] in Id.lhs
 844                              mkSysLocal (mkFastString str) uniq poly_ty
 845                            where
 846                              str     = "poly_" ++ occNameString (getOccName bndr)
 847                              poly_ty = mkPiTypes abs_vars (idType bndr)
 848
 849 newLvlVar :: String
 850           -> [CoreBndr] -> Type         -- Abstract wrt these bndrs
 851           -> LvlM Id
 852 newLvlVar str vars body_ty = do
 853     uniq <- getUniqueM
 854     return (mkSysLocal (mkFastString str) uniq (mkPiTypes vars body_ty))
 855
 856 -- The deeply tiresome thing is that we have to apply the substitution
 857 -- to the rules inside each Id.  Grr.  But it matters.
 858
 859 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
 860 cloneVar TopLevel env v ctxt_lvl dest_lvl
 861   = return (env, v)     -- Don't clone top level things
 862 cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
 863   = ASSERT( isId v ) do
 864     us <- getUniqueSupplyM
 865     let
 866       (subst', v1) = cloneIdBndr subst us v
 867       v2           = zap_demand ctxt_lvl dest_lvl v1
 868       env'         = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
 869     return (env', v2)
 870
 871 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
 872 cloneRecVars TopLevel env vs ctxt_lvl dest_lvl
 873   = return (env, vs)    -- Don't clone top level things
 874 cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
 875   = ASSERT( all isId vs ) do
 876     us <- getUniqueSupplyM
 877     let
 878       (subst', vs1) = cloneRecIdBndrs subst us vs
 879       vs2           = map (zap_demand ctxt_lvl dest_lvl) vs1
 880       env'          = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
 881     return (env', vs2)
 882
 883         -- VERY IMPORTANT: we must zap the demand info
 884         -- if the thing is going to float out past a lambda,
 885         -- or if it's going to top level (where things can't be strict)
 886 zap_demand dest_lvl ctxt_lvl id
 887   | ctxt_lvl == dest_lvl,
 888     not (isTopLvl dest_lvl) = id        -- Stays, and not going to top level
 889   | otherwise               = zapDemandIdInfo id        -- Floats out
 890 \end{code}
 891