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