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