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