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