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