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