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
 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 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
 514         -- Find the variables in fvs, free vars of the target expresion,
 515         -- whose level is less than than the supplied level
 516         -- These are the ones we are going to abstract out
 517 abstractVars dest_lvl env fvs
 518   = uniq (sortLt lt [var | fv <- varSetElems fvs, var <- absVarsOf dest_lvl env fv])
 519   where
 520         -- Sort the variables so we don't get
 521         -- mixed-up tyvars and Ids; it's just messy
 522     v1 `lt` v2 = case (isId v1, isId v2) of
 523                    (True, False) -> False
 524                    (False, True) -> True
 525                    other         -> v1 < v2     -- Same family
 526     uniq :: [Var] -> [Var]
 527         -- Remove adjacent duplicates; the sort will have brought them together
 528     uniq (v1:v2:vs) | v1 == v2  = uniq (v2:vs)
 529                     | otherwise = v1 : uniq (v2:vs)
 530     uniq vs = vs
 531
 532   -- Destintion level is the max Id level of the expression
 533   -- (We'll abstract the type variables, if any.)
 534 destLevel :: LevelEnv -> VarSet -> Bool -> Level
 535 destLevel env fvs is_function
 536   |  floatLams env
 537   && is_function = tOP_LEVEL            -- Send functions to top level; see
 538                                         -- the comments with isFunction
 539   | otherwise    = maxIdLevel env fvs
 540
 541 isFunction :: CoreExprWithFVs -> Bool
 542 -- The idea here is that we want to float *functions* to
 543 -- the top level.  This saves no work, but
 544 --      (a) it can make the host function body a lot smaller,
 545 --              and hence inlinable.
 546 --      (b) it can also save allocation when the function is recursive:
 547 --          h = \x -> letrec f = \y -> ...f...y...x...
 548 --                    in f x
 549 --     becomes
 550 --          f = \x y -> ...(f x)...y...x...
 551 --          h = \x -> f x x
 552 --     No allocation for f now.
 553 -- We may only want to do this if there are sufficiently few free
 554 -- variables.  We certainly only want to do it for values, and not for
 555 -- constructors.  So the simple thing is just to look for lambdas
 556 isFunction (_, AnnLam b e) | isId b    = True
 557                            | otherwise = isFunction e
 558 isFunction (_, AnnNote n e)            = isFunction e
 559 isFunction other                       = False
 560 \end{code}
 561
 562
 563 %************************************************************************
 564 %*                                                                      *
 565 \subsection{Free-To-Level Monad}
 566 %*                                                                      *
 567 %************************************************************************
 568
 569 \begin{code}
 570 type LevelEnv = (Bool,                          -- True <=> Float lambdas too
 571                  VarEnv Level,                  -- Domain is *post-cloned* TyVars and Ids
 572                  Subst,                         -- Domain is pre-cloned Ids; tracks the in-scope set
 573                                                 --      so that subtitution is capture-avoiding
 574                  IdEnv ([Var], LevelledExpr))   -- Domain is pre-cloned Ids
 575         -- We clone let-bound variables so that they are still
 576         -- distinct when floated out; hence the SubstEnv/IdEnv.
 577         -- (see point 3 of the module overview comment).
 578         -- We also use these envs when making a variable polymorphic
 579         -- because we want to float it out past a big lambda.
 580         --
 581         -- The SubstEnv and IdEnv always implement the same mapping, but the
 582         -- SubstEnv maps to CoreExpr and the IdEnv to LevelledExpr
 583         -- Since the range is always a variable or type application,
 584         -- there is never any difference between the two, but sadly
 585         -- the types differ.  The SubstEnv is used when substituting in
 586         -- a variable's IdInfo; the IdEnv when we find a Var.
 587         --
 588         -- In addition the IdEnv records a list of tyvars free in the
 589         -- type application, just so we don't have to call freeVars on
 590         -- the type application repeatedly.
 591         --
 592         -- The domain of the both envs is *pre-cloned* Ids, though
 593         --
 594         -- The domain of the VarEnv Level is the *post-cloned* Ids
 595
 596 initialEnv :: Bool -> LevelEnv
 597 initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
 598
 599 floatLams :: LevelEnv -> Bool
 600 floatLams (float_lams, _, _, _) = float_lams
 601
 602 extendLvlEnv :: LevelEnv -> [(Var,Level)] -> LevelEnv
 603 -- Used when *not* cloning
 604 extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
 605   = (float_lams,
 606      foldl add_lvl lvl_env prs,
 607      foldl del_subst subst prs,
 608      foldl del_id id_env prs)
 609   where
 610     add_lvl   env (v,l) = extendVarEnv env v l
 611     del_subst env (v,_) = extendInScope env v
 612     del_id    env (v,_) = delVarEnv env v
 613   -- We must remove any clone for this variable name in case of
 614   -- shadowing.  This bit me in the following case
 615   -- (in nofib/real/gg/Spark.hs):
 616   --
 617   --   case ds of wild {
 618   --     ... -> case e of wild {
 619   --              ... -> ... wild ...
 620   --            }
 621   --   }
 622   --
 623   -- The inside occurrence of @wild@ was being replaced with @ds@,
 624   -- incorrectly, because the SubstEnv was still lying around.  Ouch!
 625   -- KSW 2000-07.
 626
 627 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
 628 -- (see point 4 of the module overview comment)
 629 extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
 630   = (float_lams,
 631      extendVarEnv lvl_env case_bndr lvl,
 632      extendSubst subst case_bndr (DoneEx (Var scrut_var)),
 633      extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
 634
 635 extendCaseBndrLvlEnv env scrut case_bndr lvl
 636   = extendLvlEnv          env [(case_bndr,lvl)]
 637
 638 extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
 639   = (float_lams,
 640      foldl add_lvl   lvl_env bndr_pairs,
 641      foldl add_subst subst   bndr_pairs,
 642      foldl add_id    id_env  bndr_pairs)
 643   where
 644      add_lvl   env (v,v') = extendVarEnv env v' dest_lvl
 645      add_subst env (v,v') = extendSubst  env v (DoneEx (mkVarApps (Var v') abs_vars))
 646      add_id    env (v,v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
 647
 648 extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
 649   = (float_lams,
 650      foldl add_lvl   lvl_env bndr_pairs,
 651      new_subst,
 652      foldl add_id    id_env  bndr_pairs)
 653   where
 654      add_lvl   env (v,v') = extendVarEnv env v' lvl
 655      add_id    env (v,v') = extendVarEnv env v ([v'], Var v')
 656
 657
 658 maxIdLevel :: LevelEnv -> VarSet -> Level
 659 maxIdLevel (_, lvl_env,_,id_env) var_set
 660   = foldVarSet max_in tOP_LEVEL var_set
 661   where
 662     max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
 663                                                 Just (abs_vars, _) -> abs_vars
 664                                                 Nothing            -> [in_var])
 665
 666     max_out out_var lvl
 667         | isId out_var = case lookupVarEnv lvl_env out_var of
 668                                 Just lvl' -> maxLvl lvl' lvl
 669                                 Nothing   -> lvl
 670         | otherwise    = lvl    -- Ignore tyvars in *maxIdLevel*
 671
 672 lookupVar :: LevelEnv -> Id -> LevelledExpr
 673 lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
 674                                        Just (_, expr) -> expr
 675                                        other          -> Var v
 676
 677 absVarsOf :: Level -> LevelEnv -> Var -> [Var]
 678         -- If f is free in the exression, and f maps to poly_f a b c in the
 679         -- current substitution, then we must report a b c as candidate type
 680         -- variables
 681 absVarsOf dest_lvl (_, lvl_env, _, id_env) v
 682   | isId v
 683   = [final_av | av <- lookup_avs v, abstract_me av, final_av <- add_tyvars av]
 684
 685   | otherwise
 686   = if abstract_me v then [v] else []
 687
 688   where
 689     abstract_me v = case lookupVarEnv lvl_env v of
 690                         Just lvl -> dest_lvl `ltLvl` lvl
 691                         Nothing  -> False
 692
 693     lookup_avs v = case lookupVarEnv id_env v of
 694                         Just (abs_vars, _) -> abs_vars
 695                         Nothing            -> [v]
 696
 697         -- We are going to lambda-abstract, so nuke any IdInfo,
 698         -- and add the tyvars of the Id
 699     add_tyvars v | isId v    =  zap v  : varSetElems (idFreeTyVars v)
 700                  | otherwise = [v]
 701
 702     zap v = WARN( workerExists (idWorkerInfo v)
 703                   || not (isEmptyCoreRules (idSpecialisation v)),
 704                   text "absVarsOf: discarding info on" <+> ppr v )
 705             setIdInfo v vanillaIdInfo
 706 \end{code}
 707
 708 \begin{code}
 709 type LvlM result = UniqSM result
 710
 711 initLvl         = initUs_
 712 thenLvl         = thenUs
 713 returnLvl       = returnUs
 714 mapLvl          = mapUs
 715 \end{code}
 716
 717 \begin{code}
 718 newPolyBndrs dest_lvl env abs_vars bndrs
 719   = getUniquesUs (length bndrs)         `thenLvl` \ uniqs ->
 720     let
 721         new_bndrs = zipWith mk_poly_bndr bndrs uniqs
 722     in
 723     returnLvl (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
 724   where
 725     mk_poly_bndr bndr uniq = mkSysLocal (_PK_ str) uniq poly_ty
 726                            where
 727                              str     = "poly_" ++ occNameUserString (getOccName bndr)
 728                              poly_ty = foldr mkPiType (idType bndr) abs_vars
 729
 730
 731 newLvlVar :: String
 732           -> [CoreBndr] -> Type         -- Abstract wrt these bndrs
 733           -> LvlM Id
 734 newLvlVar str vars body_ty
 735   = getUniqueUs `thenLvl` \ uniq ->
 736     returnUs (mkSysLocal (_PK_ str) uniq (foldr mkPiType body_ty vars))
 737
 738 -- The deeply tiresome thing is that we have to apply the substitution
 739 -- to the rules inside each Id.  Grr.  But it matters.
 740
 741 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
 742 cloneVar TopLevel env v ctxt_lvl dest_lvl
 743   = returnUs (env, v)   -- Don't clone top level things
 744 cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
 745   = ASSERT( isId v )
 746     getUs       `thenLvl` \ us ->
 747     let
 748       (subst', v1) = substAndCloneId subst us v
 749       v2           = zap_demand ctxt_lvl dest_lvl v1
 750       env'         = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
 751     in
 752     returnUs (env', v2)
 753
 754 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
 755 cloneRecVars TopLevel env vs ctxt_lvl dest_lvl
 756   = returnUs (env, vs)  -- Don't clone top level things
 757 cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
 758   = ASSERT( all isId vs )
 759     getUs                       `thenLvl` \ us ->
 760     let
 761       (subst', vs1) = substAndCloneRecIds subst us vs
 762       vs2           = map (zap_demand ctxt_lvl dest_lvl) vs1
 763       env'          = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
 764     in
 765     returnUs (env', vs2)
 766
 767         -- VERY IMPORTANT: we must zap the demand info
 768         -- if the thing is going to float out past a lambda
 769 zap_demand dest_lvl ctxt_lvl id
 770   | ctxt_lvl == dest_lvl = id                   -- Stays put
 771   | otherwise            = zapDemandIdInfo id   -- Floats out
 772 \end{code}
 773