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