compiler/simplCore/SetLevels.lhs

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