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