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, exprIsCheap, 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,
  65                           idSpecialisation, idWorkerInfo, setIdInfo
  66                         )
  67 import IdInfo           ( workerExists, vanillaIdInfo, isEmptySpecInfo )
  68 import Var              ( 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 InlineCtxt
 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 0 0
 154 iNLINE_CTXT = InlineCtxt
 155
 156 incMajorLvl :: Level -> Level
 157 -- For InlineCtxt we ignore any inc's; we don't want
 158 -- to do any floating at all; see notes above
 159 incMajorLvl InlineCtxt          = InlineCtxt
 160 incMajorLvl (Level major minor) = Level (major+1) 0
 161
 162 incMinorLvl :: Level -> Level
 163 incMinorLvl InlineCtxt          = InlineCtxt
 164 incMinorLvl (Level major minor) = Level major (minor+1)
 165
 166 maxLvl :: Level -> Level -> Level
 167 maxLvl InlineCtxt l2  = l2
 168 maxLvl l1  InlineCtxt = l1
 169 maxLvl l1@(Level maj1 min1) l2@(Level maj2 min2)
 170   | (maj1 > maj2) || (maj1 == maj2 && min1 > min2) = l1
 171   | otherwise                                      = l2
 172
 173 ltLvl :: Level -> Level -> Bool
 174 ltLvl any_lvl    InlineCtxt  = False
 175 ltLvl InlineCtxt (Level _ _) = True
 176 ltLvl (Level maj1 min1) (Level maj2 min2)
 177   = (maj1 < maj2) || (maj1 == maj2 && min1 < min2)
 178
 179 ltMajLvl :: Level -> Level -> Bool
 180     -- Tells if one level belongs to a difft *lambda* level to another
 181 ltMajLvl any_lvl        InlineCtxt     = False
 182 ltMajLvl InlineCtxt     (Level maj2 _) = 0 < maj2
 183 ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2
 184
 185 isTopLvl :: Level -> Bool
 186 isTopLvl (Level 0 0) = True
 187 isTopLvl other       = False
 188
 189 isInlineCtxt :: Level -> Bool
 190 isInlineCtxt InlineCtxt = True
 191 isInlineCtxt other      = False
 192
 193 instance Outputable Level where
 194   ppr InlineCtxt      = text "<INLINE>"
 195   ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]
 196
 197 instance Eq Level where
 198   InlineCtxt        == InlineCtxt        = True
 199   (Level maj1 min1) == (Level maj2 min2) = maj1==maj2 && min1==min2
 200   l1                == l2                = False
 201 \end{code}
 202
 203
 204 %************************************************************************
 205 %*                                                                      *
 206 \subsection{Main level-setting code}
 207 %*                                                                      *
 208 %************************************************************************
 209
 210 \begin{code}
 211 setLevels :: FloatOutSwitches
 212           -> [CoreBind]
 213           -> UniqSupply
 214           -> [LevelledBind]
 215
 216 setLevels float_lams binds us
 217   = initLvl us (do_them binds)
 218   where
 219     -- "do_them"'s main business is to thread the monad along
 220     -- It gives each top binding the same empty envt, because
 221     -- things unbound in the envt have level number zero implicitly
 222     do_them :: [CoreBind] -> LvlM [LevelledBind]
 223
 224     do_them [] = returnLvl []
 225     do_them (b:bs)
 226       = lvlTopBind init_env b   `thenLvl` \ (lvld_bind, _) ->
 227         do_them bs              `thenLvl` \ lvld_binds ->
 228         returnLvl (lvld_bind : lvld_binds)
 229
 230     init_env = initialEnv float_lams
 231
 232 lvlTopBind env (NonRec binder rhs)
 233   = lvlBind TopLevel tOP_LEVEL env (AnnNonRec binder (freeVars rhs))
 234                                         -- Rhs can have no free vars!
 235
 236 lvlTopBind env (Rec pairs)
 237   = lvlBind TopLevel tOP_LEVEL env (AnnRec [(b,freeVars rhs) | (b,rhs) <- pairs])
 238 \end{code}
 239
 240 %************************************************************************
 241 %*                                                                      *
 242 \subsection{Setting expression levels}
 243 %*                                                                      *
 244 %************************************************************************
 245
 246 \begin{code}
 247 lvlExpr :: Level                -- ctxt_lvl: Level of enclosing expression
 248         -> LevelEnv             -- Level of in-scope names/tyvars
 249         -> CoreExprWithFVs      -- input expression
 250         -> LvlM LevelledExpr    -- Result expression
 251 \end{code}
 252
 253 The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
 254 binder.  Here's an example
 255
 256         v = \x -> ...\y -> let r = case (..x..) of
 257                                         ..x..
 258                            in ..
 259
 260 When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
 261 the level of @r@, even though it's inside a level-2 @\y@.  It's
 262 important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
 263 don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
 264 --- because it isn't a *maximal* free expression.
 265
 266 If there were another lambda in @r@'s rhs, it would get level-2 as well.
 267
 268 \begin{code}
 269 lvlExpr _ _ (_, AnnType ty)   = returnLvl (Type ty)
 270 lvlExpr _ env (_, AnnVar v)   = returnLvl (lookupVar env v)
 271 lvlExpr _ env (_, AnnLit lit) = returnLvl (Lit lit)
 272
 273 lvlExpr ctxt_lvl env (_, AnnApp fun arg)
 274   = lvl_fun fun                         `thenLvl` \ fun' ->
 275     lvlMFE  False ctxt_lvl env arg      `thenLvl` \ arg' ->
 276     returnLvl (App fun' arg')
 277   where
 278 -- gaw 2004
 279     lvl_fun (_, AnnCase _ _ _ _) = lvlMFE True ctxt_lvl env fun
 280     lvl_fun other              = lvlExpr ctxt_lvl env fun
 281         -- We don't do MFE on partial applications generally,
 282         -- but we do if the function is big and hairy, like a case
 283
 284 lvlExpr ctxt_lvl env (_, AnnNote InlineMe expr)
 285 -- Don't float anything out of an InlineMe; hence the iNLINE_CTXT
 286   = lvlExpr iNLINE_CTXT env expr        `thenLvl` \ expr' ->
 287     returnLvl (Note InlineMe expr')
 288
 289 lvlExpr ctxt_lvl env (_, AnnNote note expr)
 290   = lvlExpr ctxt_lvl env expr           `thenLvl` \ expr' ->
 291     returnLvl (Note note expr')
 292
 293 lvlExpr ctxt_lvl env (_, AnnCast expr co)
 294   = lvlExpr ctxt_lvl env expr           `thenLvl` \ expr' ->
 295     returnLvl (Cast expr' co)
 296
 297 -- We don't split adjacent lambdas.  That is, given
 298 --      \x y -> (x+1,y)
 299 -- we don't float to give
 300 --      \x -> let v = x+y in \y -> (v,y)
 301 -- Why not?  Because partial applications are fairly rare, and splitting
 302 -- lambdas makes them more expensive.
 303
 304 lvlExpr ctxt_lvl env expr@(_, AnnLam bndr rhs)
 305   = lvlMFE True new_lvl new_env body    `thenLvl` \ new_body ->
 306     returnLvl (mkLams new_bndrs new_body)
 307   where
 308     (bndrs, body)        = collectAnnBndrs expr
 309     (new_lvl, new_bndrs) = lvlLamBndrs ctxt_lvl bndrs
 310     new_env              = extendLvlEnv env new_bndrs
 311         -- At one time we called a special verion of collectBinders,
 312         -- which ignored coercions, because we don't want to split
 313         -- a lambda like this (\x -> coerce t (\s -> ...))
 314         -- This used to happen quite a bit in state-transformer programs,
 315         -- but not nearly so much now non-recursive newtypes are transparent.
 316         -- [See SetLevels rev 1.50 for a version with this approach.]
 317
 318 lvlExpr ctxt_lvl env (_, AnnLet (AnnNonRec bndr rhs) body)
 319   | isUnLiftedType (idType bndr)
 320         -- Treat unlifted let-bindings (let x = b in e) just like (case b of x -> e)
 321         -- That is, leave it exactly where it is
 322         -- We used to float unlifted bindings too (e.g. to get a cheap primop
 323         -- outside a lambda (to see how, look at lvlBind in rev 1.58)
 324         -- but an unrelated change meant that these unlifed bindings
 325         -- could get to the top level which is bad.  And there's not much point;
 326         -- unlifted bindings are always cheap, and so hardly worth floating.
 327   = lvlExpr ctxt_lvl env rhs            `thenLvl` \ rhs' ->
 328     lvlExpr incd_lvl env' body          `thenLvl` \ body' ->
 329     returnLvl (Let (NonRec bndr' rhs') body')
 330   where
 331     incd_lvl = incMinorLvl ctxt_lvl
 332     bndr' = TB bndr incd_lvl
 333     env'  = extendLvlEnv env [bndr']
 334
 335 lvlExpr ctxt_lvl env (_, AnnLet bind body)
 336   = lvlBind NotTopLevel ctxt_lvl env bind       `thenLvl` \ (bind', new_env) ->
 337     lvlExpr ctxt_lvl new_env body               `thenLvl` \ body' ->
 338     returnLvl (Let bind' body')
 339
 340 lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr ty alts)
 341   = lvlMFE True ctxt_lvl env expr       `thenLvl` \ expr' ->
 342     let
 343         alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
 344     in
 345     mapLvl (lvl_alt alts_env) alts      `thenLvl` \ alts' ->
 346     returnLvl (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)
 351         = lvlMFE True incd_lvl new_env rhs      `thenLvl` \ rhs' ->
 352           returnLvl (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 strict_ctxt ctxt_lvl env (_, AnnType ty)
 376   = returnLvl (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   = lvlFloatRhs abs_vars dest_lvl env ann_expr  `thenLvl` \ expr' ->
 389     newLvlVar "lvl" abs_vars ty                 `thenLvl` \ var ->
 390     returnLvl (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         = not (exprIsCheap expr) || isTopLvl dest_lvl
 403           -- Even if it escapes a value lambda, we only
 404           -- float if it's not cheap (unless it'll get all the
 405           -- way to the top).  I've seen cases where we
 406           -- float dozens of tiny free expressions, which cost
 407           -- more to allocate than to evaluate.
 408           -- NB: exprIsCheap is also true of bottom expressions, which
 409           --     is good; we don't want to share them
 410           --
 411           -- It's only Really Bad to float a cheap expression out of a
 412           -- strict context, because that builds a thunk that otherwise
 413           -- would never be built.  So another alternative would be to
 414           -- add
 415           --    || (strict_ctxt && not (exprIsBottom expr))
 416           -- to the condition above. We should really try this out.
 417
 418         | otherwise             -- Does not escape a value lambda
 419         = isTopLvl dest_lvl     -- Only float if we are going to the top level
 420         && floatConsts env      --   and the floatConsts flag is on
 421         && not strict_ctxt      -- Don't float from a strict context
 422           -- We are keen to float something to the top level, even if it does not
 423           -- escape a lambda, because then it needs no allocation.  But it's controlled
 424           -- by a flag, because doing this too early loses opportunities for RULES
 425           -- which (needless to say) are important in some nofib programs
 426           -- (gcd is an example).
 427           --
 428           -- Beware:
 429           --    concat = /\ a -> foldr ..a.. (++) []
 430           -- was getting turned into
 431           --    concat = /\ a -> lvl a
 432           --    lvl    = /\ a -> foldr ..a.. (++) []
 433           -- which is pretty stupid.  Hence the strict_ctxt test
 434 \end{code}
 435
 436
 437 %************************************************************************
 438 %*                                                                      *
 439 \subsection{Bindings}
 440 %*                                                                      *
 441 %************************************************************************
 442
 443 The binding stuff works for top level too.
 444
 445 \begin{code}
 446 lvlBind :: TopLevelFlag         -- Used solely to decide whether to clone
 447         -> Level                -- Context level; might be Top even for bindings nested in the RHS
 448                                 -- of a top level binding
 449         -> LevelEnv
 450         -> CoreBindWithFVs
 451         -> LvlM (LevelledBind, LevelEnv)
 452
 453 lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
 454   | isInlineCtxt ctxt_lvl               -- Don't do anything inside InlineMe
 455   = lvlExpr ctxt_lvl env rhs                    `thenLvl` \ rhs' ->
 456     returnLvl (NonRec (TB bndr ctxt_lvl) rhs', env)
 457
 458   | null abs_vars
 459   =     -- No type abstraction; clone existing binder
 460     lvlExpr dest_lvl env rhs                    `thenLvl` \ rhs' ->
 461     cloneVar top_lvl env bndr ctxt_lvl dest_lvl `thenLvl` \ (env', bndr') ->
 462     returnLvl (NonRec (TB bndr' dest_lvl) rhs', env')
 463
 464   | otherwise
 465   = -- Yes, type abstraction; create a new binder, extend substitution, etc
 466     lvlFloatRhs abs_vars dest_lvl env rhs       `thenLvl` \ rhs' ->
 467     newPolyBndrs dest_lvl env abs_vars [bndr]   `thenLvl` \ (env', [bndr']) ->
 468     returnLvl (NonRec (TB bndr' dest_lvl) rhs', env')
 469
 470   where
 471     bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
 472     abs_vars = abstractVars dest_lvl env bind_fvs
 473     dest_lvl = destLevel env bind_fvs (isFunction rhs)
 474 \end{code}
 475
 476
 477 \begin{code}
 478 lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
 479   | isInlineCtxt ctxt_lvl       -- Don't do anything inside InlineMe
 480   = mapLvl (lvlExpr ctxt_lvl env) rhss                  `thenLvl` \ rhss' ->
 481     returnLvl (Rec ([TB b ctxt_lvl | b <- bndrs] `zip` rhss'), env)
 482
 483   | null abs_vars
 484   = cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl    `thenLvl` \ (new_env, new_bndrs) ->
 485     mapLvl (lvlExpr ctxt_lvl new_env) rhss              `thenLvl` \ new_rhss ->
 486     returnLvl (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
 487
 488   | isSingleton pairs && count isId abs_vars > 1
 489   =     -- Special case for self recursion where there are
 490         -- several variables carried around: build a local loop:
 491         --      poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
 492         -- This just makes the closures a bit smaller.  If we don't do
 493         -- this, allocation rises significantly on some programs
 494         --
 495         -- We could elaborate it for the case where there are several
 496         -- mutually functions, but it's quite a bit more complicated
 497         --
 498         -- This all seems a bit ad hoc -- sigh
 499     let
 500         (bndr,rhs) = head pairs
 501         (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 502         rhs_env = extendLvlEnv env abs_vars_w_lvls
 503     in
 504     cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl   `thenLvl` \ (rhs_env', new_bndr) ->
 505     let
 506         (lam_bndrs, rhs_body)     = collectAnnBndrs rhs
 507         (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
 508         body_env                  = extendLvlEnv rhs_env' new_lam_bndrs
 509     in
 510     lvlExpr body_lvl body_env rhs_body          `thenLvl` \ new_rhs_body ->
 511     newPolyBndrs dest_lvl env abs_vars [bndr]   `thenLvl` \ (poly_env, [poly_bndr]) ->
 512     returnLvl (Rec [(TB poly_bndr dest_lvl,
 513                mkLams abs_vars_w_lvls $
 514                mkLams new_lam_bndrs $
 515                Let (Rec [(TB new_bndr rhs_lvl, mkLams new_lam_bndrs new_rhs_body)])
 516                    (mkVarApps (Var new_bndr) lam_bndrs))],
 517                poly_env)
 518
 519   | otherwise   -- Non-null abs_vars
 520   = newPolyBndrs dest_lvl env abs_vars bndrs            `thenLvl` \ (new_env, new_bndrs) ->
 521     mapLvl (lvlFloatRhs abs_vars dest_lvl new_env) rhss `thenLvl` \ new_rhss ->
 522     returnLvl (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
 523
 524   where
 525     (bndrs,rhss) = unzip pairs
 526
 527         -- Finding the free vars of the binding group is annoying
 528     bind_fvs        = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
 529                                     | (bndr, (rhs_fvs,_)) <- pairs])
 530                       `minusVarSet`
 531                       mkVarSet bndrs
 532
 533     dest_lvl = destLevel env bind_fvs (all isFunction rhss)
 534     abs_vars = abstractVars dest_lvl env bind_fvs
 535
 536 ----------------------------------------------------
 537 -- Three help functons for the type-abstraction case
 538
 539 lvlFloatRhs abs_vars dest_lvl env rhs
 540   = lvlExpr rhs_lvl rhs_env rhs `thenLvl` \ rhs' ->
 541     returnLvl (mkLams abs_vars_w_lvls rhs')
 542   where
 543     (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
 544     rhs_env = extendLvlEnv env abs_vars_w_lvls
 545 \end{code}
 546
 547
 548 %************************************************************************
 549 %*                                                                      *
 550 \subsection{Deciding floatability}
 551 %*                                                                      *
 552 %************************************************************************
 553
 554 \begin{code}
 555 lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [TaggedBndr Level])
 556 -- Compute the levels for the binders of a lambda group
 557 -- The binders returned are exactly the same as the ones passed,
 558 -- but they are now paired with a level
 559 lvlLamBndrs lvl []
 560   = (lvl, [])
 561
 562 lvlLamBndrs lvl bndrs
 563   = go  (incMinorLvl lvl)
 564         False   -- Havn't bumped major level in this group
 565         [] bndrs
 566   where
 567     go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
 568         | isId bndr &&                  -- Go to the next major level if this is a value binder,
 569           not bumped_major &&           -- and we havn't already gone to the next level (one jump per group)
 570           not (isOneShotLambda bndr)    -- and it isn't a one-shot lambda
 571         = go new_lvl True (TB bndr new_lvl : rev_lvld_bndrs) bndrs
 572
 573         | otherwise
 574         = go old_lvl bumped_major (TB bndr old_lvl : rev_lvld_bndrs) bndrs
 575
 576         where
 577           new_lvl = incMajorLvl old_lvl
 578
 579     go old_lvl _ rev_lvld_bndrs []
 580         = (old_lvl, reverse rev_lvld_bndrs)
 581         -- a lambda like this (\x -> coerce t (\s -> ...))
 582         -- This happens quite a bit in state-transformer programs
 583 \end{code}
 584
 585 \begin{code}
 586   -- Destintion level is the max Id level of the expression
 587   -- (We'll abstract the type variables, if any.)
 588 destLevel :: LevelEnv -> VarSet -> Bool -> Level
 589 destLevel env fvs is_function
 590   |  floatLams env
 591   && is_function = tOP_LEVEL            -- Send functions to top level; see
 592                                         -- the comments with isFunction
 593   | otherwise    = maxIdLevel env fvs
 594
 595 isFunction :: CoreExprWithFVs -> Bool
 596 -- The idea here is that we want to float *functions* to
 597 -- the top level.  This saves no work, but
 598 --      (a) it can make the host function body a lot smaller,
 599 --              and hence inlinable.
 600 --      (b) it can also save allocation when the function is recursive:
 601 --          h = \x -> letrec f = \y -> ...f...y...x...
 602 --                    in f x
 603 --     becomes
 604 --          f = \x y -> ...(f x)...y...x...
 605 --          h = \x -> f x x
 606 --     No allocation for f now.
 607 -- We may only want to do this if there are sufficiently few free
 608 -- variables.  We certainly only want to do it for values, and not for
 609 -- constructors.  So the simple thing is just to look for lambdas
 610 isFunction (_, AnnLam b e) | isId b    = True
 611                            | otherwise = isFunction e
 612 isFunction (_, AnnNote n e)            = isFunction e
 613 isFunction other                       = False
 614 \end{code}
 615
 616
 617 %************************************************************************
 618 %*                                                                      *
 619 \subsection{Free-To-Level Monad}
 620 %*                                                                      *
 621 %************************************************************************
 622
 623 \begin{code}
 624 type LevelEnv = (FloatOutSwitches,
 625                  VarEnv Level,                  -- Domain is *post-cloned* TyVars and Ids
 626                  Subst,                         -- Domain is pre-cloned Ids; tracks the in-scope set
 627                                                 --      so that subtitution is capture-avoiding
 628                  IdEnv ([Var], LevelledExpr))   -- Domain is pre-cloned Ids
 629         -- We clone let-bound variables so that they are still
 630         -- distinct when floated out; hence the SubstEnv/IdEnv.
 631         -- (see point 3 of the module overview comment).
 632         -- We also use these envs when making a variable polymorphic
 633         -- because we want to float it out past a big lambda.
 634         --
 635         -- The SubstEnv and IdEnv always implement the same mapping, but the
 636         -- SubstEnv maps to CoreExpr and the IdEnv to LevelledExpr
 637         -- Since the range is always a variable or type application,
 638         -- there is never any difference between the two, but sadly
 639         -- the types differ.  The SubstEnv is used when substituting in
 640         -- a variable's IdInfo; the IdEnv when we find a Var.
 641         --
 642         -- In addition the IdEnv records a list of tyvars free in the
 643         -- type application, just so we don't have to call freeVars on
 644         -- the type application repeatedly.
 645         --
 646         -- The domain of the both envs is *pre-cloned* Ids, though
 647         --
 648         -- The domain of the VarEnv Level is the *post-cloned* Ids
 649
 650 initialEnv :: FloatOutSwitches -> LevelEnv
 651 initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
 652
 653 floatLams :: LevelEnv -> Bool
 654 floatLams (FloatOutSw float_lams _, _, _, _) = float_lams
 655
 656 floatConsts :: LevelEnv -> Bool
 657 floatConsts (FloatOutSw _ float_consts, _, _, _) = float_consts
 658
 659 extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
 660 -- Used when *not* cloning
 661 extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
 662   = (float_lams,
 663      foldl add_lvl lvl_env prs,
 664      foldl del_subst subst prs,
 665      foldl del_id id_env prs)
 666   where
 667     add_lvl   env (TB v l) = extendVarEnv env v l
 668     del_subst env (TB v _) = extendInScope env v
 669     del_id    env (TB v _) = delVarEnv env v
 670   -- We must remove any clone for this variable name in case of
 671   -- shadowing.  This bit me in the following case
 672   -- (in nofib/real/gg/Spark.hs):
 673   --
 674   --   case ds of wild {
 675   --     ... -> case e of wild {
 676   --              ... -> ... wild ...
 677   --            }
 678   --   }
 679   --
 680   -- The inside occurrence of @wild@ was being replaced with @ds@,
 681   -- incorrectly, because the SubstEnv was still lying around.  Ouch!
 682   -- KSW 2000-07.
 683
 684 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
 685 -- (see point 4 of the module overview comment)
 686 extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
 687   = (float_lams,
 688      extendVarEnv lvl_env case_bndr lvl,
 689      extendIdSubst subst case_bndr (Var scrut_var),
 690      extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
 691
 692 extendCaseBndrLvlEnv env scrut case_bndr lvl
 693   = extendLvlEnv          env [TB case_bndr lvl]
 694
 695 extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
 696   = (float_lams,
 697      foldl add_lvl   lvl_env bndr_pairs,
 698      foldl add_subst subst   bndr_pairs,
 699      foldl add_id    id_env  bndr_pairs)
 700   where
 701      add_lvl   env (v,v') = extendVarEnv env v' dest_lvl
 702      add_subst env (v,v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
 703      add_id    env (v,v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
 704
 705 extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
 706   = (float_lams,
 707      foldl add_lvl   lvl_env bndr_pairs,
 708      new_subst,
 709      foldl add_id    id_env  bndr_pairs)
 710   where
 711      add_lvl   env (v,v') = extendVarEnv env v' lvl
 712      add_id    env (v,v') = extendVarEnv env v ([v'], Var v')
 713
 714
 715 maxIdLevel :: LevelEnv -> VarSet -> Level
 716 maxIdLevel (_, lvl_env,_,id_env) var_set
 717   = foldVarSet max_in tOP_LEVEL var_set
 718   where
 719     max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
 720                                                 Just (abs_vars, _) -> abs_vars
 721                                                 Nothing            -> [in_var])
 722
 723     max_out out_var lvl
 724         | isId out_var = case lookupVarEnv lvl_env out_var of
 725                                 Just lvl' -> maxLvl lvl' lvl
 726                                 Nothing   -> lvl
 727         | otherwise    = lvl    -- Ignore tyvars in *maxIdLevel*
 728
 729 lookupVar :: LevelEnv -> Id -> LevelledExpr
 730 lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
 731                                        Just (_, expr) -> expr
 732                                        other          -> Var v
 733
 734 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
 735         -- Find the variables in fvs, free vars of the target expresion,
 736         -- whose level is greater than the destination level
 737         -- These are the ones we are going to abstract out
 738 abstractVars dest_lvl env fvs
 739   = uniq (sortLe le [var | fv <- varSetElems fvs, var <- absVarsOf dest_lvl env fv])
 740   where
 741         -- Sort the variables so we don't get
 742         -- mixed-up tyvars and Ids; it's just messy
 743     v1 `le` v2 = case (isId v1, isId v2) of
 744                    (True, False) -> False
 745                    (False, True) -> True
 746                    other         -> v1 <= v2    -- Same family
 747
 748     uniq :: [Var] -> [Var]
 749         -- Remove adjacent duplicates; the sort will have brought them together
 750     uniq (v1:v2:vs) | v1 == v2  = uniq (v2:vs)
 751                     | otherwise = v1 : uniq (v2:vs)
 752     uniq vs = vs
 753
 754 absVarsOf :: Level -> LevelEnv -> Var -> [Var]
 755         -- If f is free in the expression, and f maps to poly_f a b c in the
 756         -- current substitution, then we must report a b c as candidate type
 757         -- variables
 758 absVarsOf dest_lvl (_, lvl_env, _, id_env) v
 759   | isId v
 760   = [zap av2 | av1 <- lookup_avs v, av2 <- add_tyvars av1, abstract_me av2]
 761
 762   | otherwise
 763   = if abstract_me v then [v] else []
 764
 765   where
 766     abstract_me v = case lookupVarEnv lvl_env v of
 767                         Just lvl -> dest_lvl `ltLvl` lvl
 768                         Nothing  -> False
 769
 770     lookup_avs v = case lookupVarEnv id_env v of
 771                         Just (abs_vars, _) -> abs_vars
 772                         Nothing            -> [v]
 773
 774     add_tyvars v = v : varSetElems (varTypeTyVars v)
 775
 776         -- We are going to lambda-abstract, so nuke any IdInfo,
 777         -- and add the tyvars of the Id (if necessary)
 778     zap v | isId v = WARN( workerExists (idWorkerInfo v) ||
 779                            not (isEmptySpecInfo (idSpecialisation v)),
 780                            text "absVarsOf: discarding info on" <+> ppr v )
 781                      setIdInfo v vanillaIdInfo
 782           | otherwise = v
 783 \end{code}
 784
 785 \begin{code}
 786 type LvlM result = UniqSM result
 787
 788 initLvl         = initUs_
 789 thenLvl         = thenUs
 790 returnLvl       = returnUs
 791 mapLvl          = mapUs
 792 \end{code}
 793
 794 \begin{code}
 795 newPolyBndrs dest_lvl env abs_vars bndrs
 796   = getUniquesUs                `thenLvl` \ uniqs ->
 797     let
 798         new_bndrs = zipWith mk_poly_bndr bndrs uniqs
 799     in
 800     returnLvl (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
 801   where
 802     mk_poly_bndr bndr uniq = mkSysLocal (mkFastString str) uniq poly_ty
 803                            where
 804                              str     = "poly_" ++ occNameString (getOccName bndr)
 805                              poly_ty = mkPiTypes abs_vars (idType bndr)
 806
 807
 808 newLvlVar :: String
 809           -> [CoreBndr] -> Type         -- Abstract wrt these bndrs
 810           -> LvlM Id
 811 newLvlVar str vars body_ty
 812   = getUniqueUs `thenLvl` \ uniq ->
 813     returnUs (mkSysLocal (mkFastString str) uniq (mkPiTypes vars body_ty))
 814
 815 -- The deeply tiresome thing is that we have to apply the substitution
 816 -- to the rules inside each Id.  Grr.  But it matters.
 817
 818 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
 819 cloneVar TopLevel env v ctxt_lvl dest_lvl
 820   = returnUs (env, v)   -- Don't clone top level things
 821 cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
 822   = ASSERT( isId v )
 823     getUs       `thenLvl` \ us ->
 824     let
 825       (subst', v1) = cloneIdBndr subst us v
 826       v2           = zap_demand ctxt_lvl dest_lvl v1
 827       env'         = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
 828     in
 829     returnUs (env', v2)
 830
 831 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
 832 cloneRecVars TopLevel env vs ctxt_lvl dest_lvl
 833   = returnUs (env, vs)  -- Don't clone top level things
 834 cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
 835   = ASSERT( all isId vs )
 836     getUs                       `thenLvl` \ us ->
 837     let
 838       (subst', vs1) = cloneRecIdBndrs subst us vs
 839       vs2           = map (zap_demand ctxt_lvl dest_lvl) vs1
 840       env'          = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
 841     in
 842     returnUs (env', vs2)
 843
 844         -- VERY IMPORTANT: we must zap the demand info
 845         -- if the thing is going to float out past a lambda
 846 zap_demand dest_lvl ctxt_lvl id
 847   | ctxt_lvl == dest_lvl = id                   -- Stays put
 848   | otherwise            = zapDemandIdInfo id   -- Floats out
 849 \end{code}
 850