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
  26
  27 \begin{code}
  28 module SetLevels (
  29         setLevels,
  30
  31         Level(..), tOP_LEVEL,
  32
  33         incMinorLvl, ltMajLvl, ltLvl, isTopLvl
  34     ) where
  35
  36 #include "HsVersions.h"
  37
  38 import CoreSyn
  39
  40 import CoreUtils        ( coreExprType, exprIsTrivial, exprIsBottom )
  41 import CoreFVs          -- all of it
  42 import Id               ( Id, idType, mkSysLocal, isOneShotLambda, modifyIdInfo )
  43 import IdInfo           ( specInfo, setSpecInfo )
  44 import Var              ( IdOrTyVar, Var, setVarUnique )
  45 import VarEnv
  46 import Subst
  47 import VarSet
  48 import Type             ( isUnLiftedType, mkTyVarTys, mkForAllTys, Type )
  49 import BasicTypes       ( TopLevelFlag(..) )
  50 import VarSet
  51 import VarEnv
  52 import UniqSupply
  53 import Maybes           ( maybeToBool )
  54 import Util             ( zipWithEqual, zipEqual )
  55 import Outputable
  56
  57 isLeakFreeType x y = False -- safe option; ToDo
  58 \end{code}
  59
  60 %************************************************************************
  61 %*                                                                      *
  62 \subsection{Level numbers}
  63 %*                                                                      *
  64 %************************************************************************
  65
  66 \begin{code}
  67 data Level
  68   = Top         -- Means *really* the top level; short for (Level 0 0).
  69   | Level   Int -- Level number of enclosing lambdas
  70             Int -- Number of big-lambda and/or case expressions between
  71                 -- here and the nearest enclosing lambda
  72 \end{code}
  73
  74 The {\em level number} on a (type-)lambda-bound variable is the
  75 nesting depth of the (type-)lambda which binds it.  The outermost lambda
  76 has level 1, so (Level 0 0) means that the variable is bound outside any lambda.
  77
  78 On an expression, it's the maximum level number of its free
  79 (type-)variables.  On a let(rec)-bound variable, it's the level of its
  80 RHS.  On a case-bound variable, it's the number of enclosing lambdas.
  81
  82 Top-level variables: level~0.  Those bound on the RHS of a top-level
  83 definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
  84 as ``subscripts'')...
  85 \begin{verbatim}
  86 a_0 = let  b_? = ...  in
  87            x_1 = ... b ... in ...
  88 \end{verbatim}
  89
  90 Level 0 0 will make something get floated to a top-level "equals",
  91 @Top@ makes it go right to the top.
  92
  93 The main function @lvlExpr@ carries a ``context level'' (@ctxt_lvl@).
  94 That's meant to be the level number of the enclosing binder in the
  95 final (floated) program.  If the level number of a sub-expression is
  96 less than that of the context, then it might be worth let-binding the
  97 sub-expression so that it will indeed float. This context level starts
  98 at @Level 0 0@; it is never @Top@.
  99
 100 \begin{code}
 101 type LevelledExpr  = TaggedExpr Level
 102 type LevelledArg   = TaggedArg  Level
 103 type LevelledBind  = TaggedBind Level
 104
 105 tOP_LEVEL = Top
 106
 107 incMajorLvl :: Level -> Level
 108 incMajorLvl Top                 = Level 1 0
 109 incMajorLvl (Level major minor) = Level (major+1) 0
 110
 111 incMinorLvl :: Level -> Level
 112 incMinorLvl Top                 = Level 0 1
 113 incMinorLvl (Level major minor) = Level major (minor+1)
 114
 115 unTopify :: Type -> Level -> Level
 116 unTopify ty lvl
 117    | isUnLiftedType ty = case lvl of
 118                                 Top   -> Level 0 0      -- Unboxed floats can't go right
 119                                 other -> lvl            -- to the top
 120    | otherwise         = lvl
 121
 122 maxLvl :: Level -> Level -> Level
 123 maxLvl Top l2 = l2
 124 maxLvl l1 Top = l1
 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 l1                Top               = False
 131 ltLvl Top               (Level _ _)       = True
 132 ltLvl (Level maj1 min1) (Level maj2 min2)
 133   = (maj1 < maj2) || (maj1 == maj2 && min1 < min2)
 134
 135 ltMajLvl :: Level -> Level -> Bool
 136     -- Tells if one level belongs to a difft *lambda* level to another
 137 ltMajLvl l1             Top            = False
 138 ltMajLvl Top            (Level 0 _)    = False
 139 ltMajLvl Top            (Level _ _)    = True
 140 ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2
 141
 142 isTopLvl :: Level -> Bool
 143 isTopLvl Top   = True
 144 isTopLvl other = False
 145
 146 isTopMajLvl :: Level -> Bool -- Tells if it's the top *lambda* level
 147 isTopMajLvl Top           = True
 148 isTopMajLvl (Level maj _) = maj == 0
 149
 150 instance Outputable Level where
 151   ppr Top             = ptext SLIT("<Top>")
 152   ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]
 153 \end{code}
 154
 155 %************************************************************************
 156 %*                                                                      *
 157 \subsection{Main level-setting code}
 158 %*                                                                      *
 159 %************************************************************************
 160
 161 \begin{code}
 162 setLevels :: [CoreBind]
 163           -> UniqSupply
 164           -> [LevelledBind]
 165
 166 setLevels binds us
 167   = initLvl us (do_them binds)
 168   where
 169     -- "do_them"'s main business is to thread the monad along
 170     -- It gives each top binding the same empty envt, because
 171     -- things unbound in the envt have level number zero implicitly
 172     do_them :: [CoreBind] -> LvlM [LevelledBind]
 173
 174     do_them [] = returnLvl []
 175     do_them (b:bs)
 176       = lvlTopBind b    `thenLvl` \ (lvld_bind, _) ->
 177         do_them bs      `thenLvl` \ lvld_binds ->
 178         returnLvl (lvld_bind ++ lvld_binds)
 179
 180 lvlTopBind (NonRec binder rhs)
 181   = lvlBind TopLevel Top initialEnv (AnnNonRec binder (freeVars rhs))
 182                                         -- Rhs can have no free vars!
 183
 184 lvlTopBind (Rec pairs)
 185   = lvlBind TopLevel Top initialEnv (AnnRec [(b,freeVars rhs) | (b,rhs) <- pairs])
 186 \end{code}
 187
 188 %************************************************************************
 189 %*                                                                      *
 190 \subsection{Bindings}
 191 %*                                                                      *
 192 %************************************************************************
 193
 194 The binding stuff works for top level too.
 195
 196 \begin{code}
 197 lvlBind :: TopLevelFlag         -- Used solely to decide whether to clone
 198         -> Level                -- Context level; might be Top even for bindings nested in the RHS
 199                                 -- of a top level binding
 200         -> LevelEnv
 201         -> CoreBindWithFVs
 202         -> LvlM ([LevelledBind], LevelEnv)
 203
 204 lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs)
 205   = setFloatLevel (Just bndr) ctxt_lvl env rhs ty       `thenLvl` \ (final_lvl, rhs') ->
 206     cloneVar top_lvl env bndr final_lvl         `thenLvl` \ (new_env, new_bndr) ->
 207     returnLvl ([NonRec (new_bndr, final_lvl) rhs'], new_env)
 208   where
 209     ty = idType bndr
 210
 211
 212 lvlBind top_lvl ctxt_lvl env (AnnRec pairs) = lvlRecBind top_lvl ctxt_lvl env pairs
 213 \end{code}
 214
 215 %************************************************************************
 216 %*                                                                      *
 217 \subsection{Setting expression levels}
 218 %*                                                                      *
 219 %************************************************************************
 220
 221 \begin{code}
 222 lvlExpr :: Level                -- ctxt_lvl: Level of enclosing expression
 223         -> LevelEnv             -- Level of in-scope names/tyvars
 224         -> CoreExprWithFVs      -- input expression
 225         -> LvlM LevelledExpr    -- Result expression
 226 \end{code}
 227
 228 The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
 229 binder.
 230
 231 Here's an example
 232
 233         v = \x -> ...\y -> let r = case (..x..) of
 234                                         ..x..
 235                            in ..
 236
 237 When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
 238 the level of @r@, even though it's inside a level-2 @\y@.  It's
 239 important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
 240 don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
 241 --- because it isn't a *maximal* free expression.
 242
 243 If there were another lambda in @r@'s rhs, it would get level-2 as well.
 244
 245 \begin{code}
 246 lvlExpr _ _ (_, AnnType ty) = returnLvl (Type ty)
 247 lvlExpr _ env (_, AnnVar v) = returnLvl (lookupVar env v)
 248
 249 lvlExpr ctxt_lvl env (_, AnnCon con args)
 250   = mapLvl (lvlExpr ctxt_lvl env) args  `thenLvl` \ args' ->
 251     returnLvl (Con con args')
 252
 253 lvlExpr ctxt_lvl env (_, AnnApp fun arg)
 254   = lvlExpr ctxt_lvl env fun            `thenLvl` \ fun' ->
 255     lvlMFE  ctxt_lvl env arg            `thenLvl` \ arg' ->
 256     returnLvl (App fun' arg')
 257
 258 lvlExpr ctxt_lvl env (_, AnnNote note expr)
 259   = lvlExpr ctxt_lvl env expr           `thenLvl` \ expr' ->
 260     returnLvl (Note note expr')
 261
 262 -- We don't split adjacent lambdas.  That is, given
 263 --      \x y -> (x+1,y)
 264 -- we don't float to give
 265 --      \x -> let v = x+y in \y -> (v,y)
 266 -- Why not?  Because partial applications are fairly rare, and splitting
 267 -- lambdas makes them more expensive.
 268
 269 lvlExpr ctxt_lvl env expr@(_, AnnLam bndr rhs)
 270   = lvlMFE incd_lvl new_env body        `thenLvl` \ body' ->
 271     returnLvl (mk_lams lvld_bndrs expr body')
 272   where
 273     bndr_is_id         = isId bndr
 274     bndr_is_tyvar      = isTyVar bndr
 275     (more_bndrs, body) = go rhs
 276     bndrs              = bndr : more_bndrs
 277
 278     incd_lvl   | bndr_is_id && not (all isOneShotLambda bndrs) = incMajorLvl ctxt_lvl
 279                | otherwise                                     = incMinorLvl ctxt_lvl
 280         -- Only bump the major level number if the binders include
 281         -- at least one more-than-one-shot lambda
 282
 283     lvld_bndrs = [(b,incd_lvl) | b <- bndrs]
 284     new_env    = extendLvlEnv env lvld_bndrs
 285
 286         -- Ignore notes, because we don't want to split
 287         -- a lambda like this (\x -> coerce t (\s -> ...))
 288         -- This happens quite a bit in state-transformer programs
 289     go (_, AnnLam bndr rhs) |  bndr_is_id && isId bndr
 290                             || bndr_is_tyvar && isTyVar bndr
 291                             =  case go rhs of { (bndrs, body) -> (bndr:bndrs, body) }
 292     go (_, AnnNote _ rhs)   = go rhs
 293     go body                 = ([], body)
 294
 295         -- Have to reconstruct the right Notes, since we ignored
 296         -- them when gathering the lambdas
 297     mk_lams (lb : lbs) (_, AnnLam _ body)     body' = Lam  lb   (mk_lams lbs body body')
 298     mk_lams lbs        (_, AnnNote note body) body' = Note note (mk_lams lbs body body')
 299     mk_lams []         body                   body' = body'
 300
 301 lvlExpr ctxt_lvl env (_, AnnLet bind body)
 302   = lvlBind NotTopLevel ctxt_lvl env bind       `thenLvl` \ (binds', new_env) ->
 303     lvlExpr ctxt_lvl new_env body               `thenLvl` \ body' ->
 304     returnLvl (mkLets binds' body')
 305
 306 lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr alts)
 307   = lvlMFE ctxt_lvl env expr    `thenLvl` \ expr' ->
 308     mapLvl lvl_alt alts         `thenLvl` \ alts' ->
 309     returnLvl (Case expr' (case_bndr, incd_lvl) alts')
 310   where
 311       expr_type = coreExprType (deAnnotate expr)
 312       incd_lvl  = incMinorLvl ctxt_lvl
 313       alts_env  = extendLvlEnv env [(case_bndr,incd_lvl)]
 314
 315       lvl_alt (con, bs, rhs)
 316         = let
 317                 bs'  = [ (b, incd_lvl) | b <- bs ]
 318                 new_env = extendLvlEnv alts_env bs'
 319           in
 320           lvlMFE incd_lvl new_env rhs   `thenLvl` \ rhs' ->
 321           returnLvl (con, bs', rhs')
 322 \end{code}
 323
 324 @lvlMFE@ is just like @lvlExpr@, except that it might let-bind
 325 the expression, so that it can itself be floated.
 326
 327 \begin{code}
 328 lvlMFE ::  Level                -- Level of innermost enclosing lambda/tylam
 329         -> LevelEnv             -- Level of in-scope names/tyvars
 330         -> CoreExprWithFVs      -- input expression
 331         -> LvlM LevelledExpr    -- Result expression
 332
 333 lvlMFE ctxt_lvl env (_, AnnType ty)
 334   = returnLvl (Type ty)
 335
 336 lvlMFE ctxt_lvl env ann_expr
 337   | isUnLiftedType ty           -- Can't let-bind it
 338   = lvlExpr ctxt_lvl env ann_expr
 339
 340   | otherwise           -- Not primitive type so could be let-bound
 341   = setFloatLevel Nothing {- Not already let-bound -}
 342         ctxt_lvl env ann_expr ty        `thenLvl` \ (final_lvl, expr') ->
 343     returnLvl expr'
 344   where
 345     ty = coreExprType (deAnnotate ann_expr)
 346 \end{code}
 347
 348
 349 %************************************************************************
 350 %*                                                                      *
 351 \subsection{Deciding floatability}
 352 %*                                                                      *
 353 %************************************************************************
 354
 355 @setFloatLevel@ is used for let-bound right-hand-sides, or for MFEs which
 356 are being created as let-bindings
 357
 358 Decision tree:
 359 Let Bound?
 360   YES. -> (a) try abstracting type variables.
 361        If we abstract type variables it will go further, that is, past more
 362        lambdas. same as asking if the level number given by the free
 363        variables is less than the level number given by free variables
 364        and type variables together.
 365        Abstract offending type variables, e.g.
 366        change f ty a b
 367        to let v = /\ty' -> f ty' a b
 368           in v ty
 369        so that v' is not stopped by the level number of ty
 370        tag the original let with its level number
 371        (from its variables and type variables)
 372   NO.  is a WHNF?
 373          YES. -> No point in let binding to float a WHNF.
 374                  Pin (leave) expression here.
 375          NO. -> Will float past a lambda?
 376                 (check using free variables only, not type variables)
 377                   YES. -> do the same as (a) above.
 378                   NO. -> No point in let binding if it is not going anywhere
 379                          Pin (leave) expression here.
 380
 381 \begin{code}
 382 setFloatLevel :: Maybe Id               -- Just id <=> the expression is already let-bound to id
 383                                         -- Nothing <=> it's a possible MFE
 384               -> Level                  -- of context
 385               -> LevelEnv
 386
 387               -> CoreExprWithFVs        -- Original rhs
 388               -> Type                   -- Type of rhs
 389
 390               -> LvlM (Level,           -- Level to attribute to this let-binding
 391                        LevelledExpr)    -- Final rhs
 392
 393 setFloatLevel maybe_let_bound ctxt_lvl env expr@(expr_fvs, _) ty
 394
 395 -- Now deal with (by not floating) trivial non-let-bound expressions
 396 -- which just aren't worth let-binding in order to float.  We always
 397 -- choose to float even trivial let-bound things because it doesn't do
 398 -- any harm, and not floating it may pin something important.  For
 399 -- example
 400 --
 401 --      x = let v = []
 402 --              w = 1:v
 403 --          in ...
 404 --
 405 -- Here, if we don't float v we won't float w, which is Bad News.
 406 -- If this gives any problems we could restrict the idea to things destined
 407 -- for top level.
 408
 409   | not alreadyLetBound
 410     && (expr_is_trivial || expr_is_bottom || not will_float_past_lambda)
 411
 412   =   -- Pin trivial non-let-bound expressions,
 413       -- or ones which aren't going anywhere useful
 414     lvlExpr ctxt_lvl env expr           `thenLvl` \ expr' ->
 415     returnLvl (safe_ctxt_lvl, expr')
 416
 417 {- SDM 7/98
 418 The above case used to read (whnf_or_bottom || not will_float_past_lambda).
 419 It was changed because we really do want to float out constructors if possible:
 420 this can save a great deal of needless allocation inside a loop.  On the other
 421 hand, there's no point floating out nullary constructors and literals, hence
 422 the expr_is_trivial condition.
 423 -}
 424
 425   | alreadyLetBound && not worth_type_abstraction
 426   =   -- Process the expression with a new ctxt_lvl, obtained from
 427       -- the free vars of the expression itself
 428     lvlExpr expr_lvl env expr           `thenLvl` \ expr' ->
 429     returnLvl (safe_expr_lvl, expr')
 430
 431   | otherwise -- This will create a let anyway, even if there is no
 432               -- type variable to abstract, so we try to abstract anyway
 433   = abstractWrtTyVars offending_tyvars ty env lvl_after_ty_abstr expr
 434                                               `thenLvl` \ final_expr ->
 435     returnLvl (safe_expr_lvl, final_expr)
 436       -- OLD LIE: The body of the let, just a type application, isn't worth floating
 437       --          so pin it with ctxt_lvl
 438       -- The truth: better to give it expr_lvl in case it is pinning
 439       -- something non-trivial which depends on it.
 440   where
 441     alreadyLetBound = maybeToBool maybe_let_bound
 442
 443     safe_ctxt_lvl   = unTopify ty ctxt_lvl
 444     safe_expr_lvl   = unTopify ty expr_lvl
 445
 446     fvs                = case maybe_let_bound of
 447                                 Nothing -> expr_fvs
 448                                 Just id -> expr_fvs `unionVarSet` idFreeVars id
 449
 450     ids_only_lvl       = foldVarSet (maxIdLvl    env) tOP_LEVEL fvs
 451     tyvars_only_lvl    = foldVarSet (maxTyVarLvl env) tOP_LEVEL fvs
 452     expr_lvl           = ids_only_lvl `maxLvl` tyvars_only_lvl
 453     lvl_after_ty_abstr = ids_only_lvl --`maxLvl` non_offending_tyvars_lvl
 454
 455         -- Will escape lambda if let-bound
 456     will_float_past_lambda = ids_only_lvl `ltMajLvl` ctxt_lvl
 457
 458          -- Will escape (more) lambda(s)/type lambda(s) if type abstracted
 459     worth_type_abstraction =  (ids_only_lvl `ltLvl` tyvars_only_lvl)
 460                            && not expr_is_trivial        -- Avoids abstracting trivial type applications
 461
 462     offending_tyvars = filter offending_tv (varSetElems fvs)
 463     offending_tv var | isId var  = False
 464                      | otherwise = ids_only_lvl `ltLvl` varLevel env var
 465
 466     expr_is_trivial = exprIsTrivial de_ann_expr
 467     expr_is_bottom  = exprIsBottom  de_ann_expr
 468     de_ann_expr     = deAnnotate expr
 469 \end{code}
 470
 471 Abstract wrt tyvars, by making it just as if we had seen
 472
 473      let v = /\a1..an. E
 474      in v a1 ... an
 475
 476 instead of simply E. The idea is that v can be freely floated, since it
 477 has no free type variables. Of course, if E has no free type
 478 variables, then we just return E.
 479
 480 \begin{code}
 481 abstractWrtTyVars offending_tyvars ty env lvl expr
 482   = lvlExpr incd_lvl new_env expr       `thenLvl` \ expr' ->
 483     newLvlVar poly_ty                   `thenLvl` \ poly_var ->
 484     let
 485        poly_var_rhs     = mkLams tyvar_lvls expr'
 486        poly_var_binding = NonRec (poly_var, lvl) poly_var_rhs
 487        poly_var_app     = mkTyApps (Var poly_var) (mkTyVarTys offending_tyvars)
 488        final_expr       = Let poly_var_binding poly_var_app -- mkCoLet* requires Core
 489     in
 490     returnLvl final_expr
 491   where
 492     poly_ty = mkForAllTys offending_tyvars ty
 493
 494         -- These defns are just like those in the TyLam case of lvlExpr
 495     incd_lvl   = incMinorLvl lvl
 496     tyvar_lvls = [(tv,incd_lvl) | tv <- offending_tyvars]
 497     new_env    = extendLvlEnv env tyvar_lvls
 498 \end{code}
 499
 500 Recursive definitions.  We want to transform
 501
 502         letrec
 503            x1 = e1
 504            ...
 505            xn = en
 506         in
 507         body
 508
 509 to
 510
 511         letrec
 512            x1' = /\ ab -> let D' in e1
 513            ...
 514            xn' = /\ ab -> let D' in en
 515         in
 516         let D in body
 517
 518 where ab are the tyvars pinning the defn further in than it
 519 need be, and D is a bunch of simple type applications:
 520
 521                 x1_cl = x1' ab
 522                 ...
 523                 xn_cl = xn' ab
 524
 525 The "_cl" indicates that in D, the level numbers on the xi are the context level
 526 number; type applications aren't worth floating.  The D' decls are
 527 similar:
 528
 529                 x1_ll = x1' ab
 530                 ...
 531                 xn_ll = xn' ab
 532
 533 but differ in their level numbers; here the ab are the newly-introduced
 534 type lambdas.
 535
 536 \begin{code}
 537 lvlRecBind top_lvl ctxt_lvl env pairs
 538   | ids_only_lvl `ltLvl` tyvars_only_lvl
 539   =     -- Abstract wrt tyvars;
 540         -- offending_tyvars is definitely non-empty
 541         -- (I love the ASSERT to check this...  WDP 95/02)
 542     let
 543        incd_lvl         = incMinorLvl ids_only_lvl
 544        tyvars_w_rhs_lvl = [(var,incd_lvl) | var <- offending_tyvars]
 545        bndrs_w_rhs_lvl  = [(var,incd_lvl) | var <- bndrs]
 546        rhs_env          = extendLvlEnv env (tyvars_w_rhs_lvl ++ bndrs_w_rhs_lvl)
 547     in
 548     mapLvl (lvlExpr incd_lvl rhs_env) rhss      `thenLvl` \ rhss' ->
 549     mapLvl newLvlVar poly_tys                   `thenLvl` \ poly_vars ->
 550     cloneVars top_lvl env bndrs ctxt_lvl        `thenLvl` \ (new_env, new_bndrs) ->
 551     let
 552                 -- The "d_rhss" are the right-hand sides of "D" and "D'"
 553                 -- in the documentation above
 554         d_rhss = [ mkTyApps (Var poly_var) offending_tyvar_tys | poly_var <- poly_vars]
 555
 556                 -- "local_binds" are "D'" in the documentation above
 557         local_binds = zipWithEqual "SetLevels" NonRec bndrs_w_rhs_lvl d_rhss
 558
 559         poly_var_rhss = [ mkLams tyvars_w_rhs_lvl (mkLets local_binds rhs')
 560                         | rhs' <- rhss'
 561                         ]
 562
 563         poly_binds  = zipEqual "poly_binds" [(poly_var, ids_only_lvl) | poly_var <- poly_vars]
 564                                             poly_var_rhss
 565
 566                 -- The new right-hand sides, just a type application,
 567                 -- aren't worth floating so pin it with ctxt_lvl
 568         bndrs_w_lvl = new_bndrs `zip` repeat ctxt_lvl
 569
 570                 -- "d_binds" are the "D" in the documentation above
 571         d_binds = zipWithEqual "SetLevels" NonRec bndrs_w_lvl d_rhss
 572     in
 573     returnLvl (Rec poly_binds : d_binds, new_env)
 574
 575   | otherwise
 576   =     -- Let it float freely
 577     cloneVars top_lvl env bndrs expr_lvl                `thenLvl` \ (new_env, new_bndrs) ->
 578     let
 579         bndrs_w_lvls = new_bndrs `zip` repeat expr_lvl
 580     in
 581     mapLvl (lvlExpr expr_lvl new_env) rhss      `thenLvl` \ rhss' ->
 582     returnLvl ([Rec (bndrs_w_lvls `zip` rhss')], new_env)
 583
 584   where
 585     (bndrs,rhss) = unzip pairs
 586
 587         -- Finding the free vars of the binding group is annoying
 588     bind_fvs        = (unionVarSets (map fst rhss) `unionVarSet` unionVarSets (map idFreeVars bndrs))
 589                       `minusVarSet`
 590                       mkVarSet bndrs
 591
 592     ids_only_lvl    = foldVarSet (maxIdLvl    env) tOP_LEVEL bind_fvs
 593     tyvars_only_lvl = foldVarSet (maxTyVarLvl env) tOP_LEVEL bind_fvs
 594     expr_lvl        = ids_only_lvl `maxLvl` tyvars_only_lvl
 595
 596     offending_tyvars = filter offending_tv (varSetElems bind_fvs)
 597     offending_tv var | isId var  = False
 598                      | otherwise = ids_only_lvl `ltLvl` varLevel env var
 599     offending_tyvar_tys = mkTyVarTys offending_tyvars
 600
 601     tys      = map idType bndrs
 602     poly_tys = map (mkForAllTys offending_tyvars) tys
 603 \end{code}
 604
 605 %************************************************************************
 606 %*                                                                      *
 607 \subsection{Free-To-Level Monad}
 608 %*                                                                      *
 609 %************************************************************************
 610
 611 \begin{code}
 612 type LevelEnv = (VarEnv Level, SubstEnv)
 613         -- We clone let-bound variables so that they are still
 614         -- distinct when floated out; hence the SubstEnv
 615         -- The domain of the VarEnv is *pre-cloned* Ids, though
 616
 617 initialEnv :: LevelEnv
 618 initialEnv = (emptyVarEnv, emptySubstEnv)
 619
 620 extendLvlEnv :: LevelEnv -> [(Var,Level)] -> LevelEnv
 621         -- Used when *not* cloning
 622 extendLvlEnv (lvl_env, subst_env) prs
 623    = (foldl add lvl_env prs, subst_env)
 624    where
 625      add env (v,l) = extendVarEnv env v l
 626
 627 varLevel :: LevelEnv -> IdOrTyVar -> Level
 628 varLevel (lvl_env, _) v
 629   = case lookupVarEnv lvl_env v of
 630       Just level -> level
 631       Nothing    -> tOP_LEVEL
 632
 633 lookupVar :: LevelEnv -> Id -> LevelledExpr
 634 lookupVar (_, subst) v = case lookupSubstEnv subst v of
 635                            Just (DoneEx (Var v')) -> Var v'     -- Urgh!  Types don't match
 636                            other                  -> Var v
 637
 638 maxIdLvl :: LevelEnv -> IdOrTyVar -> Level -> Level
 639 maxIdLvl (lvl_env,_) var lvl | isTyVar var = lvl
 640                              | otherwise   = case lookupVarEnv lvl_env var of
 641                                                 Just lvl' -> maxLvl lvl' lvl
 642                                                 Nothing   -> lvl
 643
 644 maxTyVarLvl :: LevelEnv -> IdOrTyVar -> Level -> Level
 645 maxTyVarLvl (lvl_env,_) var lvl | isId var  = lvl
 646                                 | otherwise = case lookupVarEnv lvl_env var of
 647                                                 Just lvl' -> maxLvl lvl' lvl
 648                                                 Nothing   -> lvl
 649 \end{code}
 650
 651 \begin{code}
 652 type LvlM result = UniqSM result
 653
 654 initLvl         = initUs_
 655 thenLvl         = thenUs
 656 returnLvl       = returnUs
 657 mapLvl          = mapUs
 658 \end{code}
 659
 660 \begin{code}
 661 newLvlVar :: Type -> LvlM Id
 662 newLvlVar ty = getUniqueUs      `thenLvl` \ uniq ->
 663                returnUs (mkSysLocal SLIT("lvl") uniq ty)
 664
 665 -- The deeply tiresome thing is that we have to apply the substitution
 666 -- to the rules inside each Id.  Grr.  But it matters.
 667
 668 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> LvlM (LevelEnv, Id)
 669 cloneVar TopLevel env v lvl
 670   = returnUs (env, v)   -- Don't clone top level things
 671 cloneVar NotTopLevel (lvl_env, subst_env) v lvl
 672   = getUniqueUs `thenLvl` \ uniq ->
 673     let
 674       subst      = mkSubst emptyVarSet subst_env
 675       v'         = setVarUnique v uniq
 676       v''        = modifyIdInfo (\info -> substIdInfo subst info info) v'
 677       subst_env' = extendSubstEnv subst_env v (DoneEx (Var v''))
 678       lvl_env'   = extendVarEnv lvl_env v lvl
 679     in
 680     returnUs ((lvl_env', subst_env'), v'')
 681
 682 cloneVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> LvlM (LevelEnv, [Id])
 683 cloneVars TopLevel env vs lvl
 684   = returnUs (env, vs)  -- Don't clone top level things
 685 cloneVars NotTopLevel (lvl_env, subst_env) vs lvl
 686   = getUniquesUs (length vs)    `thenLvl` \ uniqs ->
 687     let
 688       subst      = mkSubst emptyVarSet subst_env'
 689       vs'        = zipWith setVarUnique vs uniqs
 690       vs''       = map (modifyIdInfo (\info -> substIdInfo subst info info)) vs'
 691       subst_env' = extendSubstEnvList subst_env vs [DoneEx (Var v'') | v'' <- vs'']
 692       lvl_env'   = extendVarEnvList lvl_env (vs `zip` repeat lvl)
 693     in
 694     returnUs ((lvl_env', subst_env'), vs'')
 695 \end{code}