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