SetLevels.lhs 30.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45
%
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
%
\section{SetLevels}

		***************************
			Overview
		***************************

1. We attach binding levels to Core bindings, in preparation for floating
   outwards (@FloatOut@).

2. We also let-ify many expressions (notably case scrutinees), so they
   will have a fighting chance of being floated sensible.

3. We clone the binders of any floatable let-binding, so that when it is
   floated out it will be unique.  (This used to be done by the simplifier
   but the latter now only ensures that there's no shadowing; indeed, even 
   that may not be true.)

   NOTE: this can't be done using the uniqAway idea, because the variable
 	 must be unique in the whole program, not just its current scope,
	 because two variables in different scopes may float out to the
	 same top level place

   NOTE: Very tiresomely, we must apply this substitution to
	 the rules stored inside a variable too.

   We do *not* clone top-level bindings, because some of them must not change,
   but we *do* clone bindings that are heading for the top level

4. In the expression
	case x of wild { p -> ...wild... }
   we substitute x for wild in the RHS of the case alternatives:
	case x of wild { p -> ...x... }
   This means that a sub-expression involving x is not "trapped" inside the RHS.
   And it's not inconvenient because we already have a substitution.

  Note that this is EXACTLY BACKWARDS from the what the simplifier does.
  The simplifier tries to get rid of occurrences of x, in favour of wild,
  in the hope that there will only be one remaining occurrence of x, namely
  the scrutinee of the case, and we can inline it.  

\begin{code}
module SetLevels (
46
	setLevels, 
47 48

	Level(..), tOP_LEVEL,
49
	LevelledBind, LevelledExpr,
50

51
	incMinorLvl, ltMajLvl, ltLvl, isTopLvl, isInlineCtxt
52 53 54 55 56 57
    ) where

#include "HsVersions.h"

import CoreSyn

58
import CmdLineOpts	( FloatOutSwitches(..) )
59
import CoreUtils	( exprType, exprIsTrivial, exprIsCheap, mkPiTypes )
60
import CoreFVs		-- all of it
61 62
import CoreSubst	( Subst, emptySubst, extendInScope, extendIdSubst,
			  cloneIdBndr, cloneRecIdBndrs )
63 64
import Id		( Id, idType, mkSysLocalUnencoded, 
			  isOneShotLambda, zapDemandIdInfo,
65 66 67 68 69 70 71 72 73 74 75
			  idSpecialisation, idWorkerInfo, setIdInfo
			)
import IdInfo		( workerExists, vanillaIdInfo, )
import Var		( Var )
import VarSet
import VarEnv
import Name		( getOccName )
import OccName		( occNameUserString )
import Type		( isUnLiftedType, Type )
import BasicTypes	( TopLevelFlag(..) )
import UniqSupply
76
import Util		( sortLe, isSingleton, count )
77
import Outputable
78
import FastString
79 80 81 82 83 84 85 86 87
\end{code}

%************************************************************************
%*									*
\subsection{Level numbers}
%*									*
%************************************************************************

\begin{code}
88 89 90
data Level = InlineCtxt	-- A level that's used only for
			-- the context parameter ctxt_lvl
	   | Level Int	-- Level number of enclosing lambdas
91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
	  	   Int	-- Number of big-lambda and/or case expressions between
			-- here and the nearest enclosing lambda
\end{code}

The {\em level number} on a (type-)lambda-bound variable is the
nesting depth of the (type-)lambda which binds it.  The outermost lambda
has level 1, so (Level 0 0) means that the variable is bound outside any lambda.

On an expression, it's the maximum level number of its free
(type-)variables.  On a let(rec)-bound variable, it's the level of its
RHS.  On a case-bound variable, it's the number of enclosing lambdas.

Top-level variables: level~0.  Those bound on the RHS of a top-level
definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
as ``subscripts'')...
\begin{verbatim}
a_0 = let  b_? = ...  in
	   x_1 = ... b ... in ...
\end{verbatim}

The main function @lvlExpr@ carries a ``context level'' (@ctxt_lvl@).
That's meant to be the level number of the enclosing binder in the
final (floated) program.  If the level number of a sub-expression is
less than that of the context, then it might be worth let-binding the
115 116 117 118
sub-expression so that it will indeed float.  

If you can float to level @Level 0 0@ worth doing so because then your
allocation becomes static instead of dynamic.  We always start with
119
context @Level 0 0@.  
120

121

122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148
InlineCtxt
~~~~~~~~~~
@InlineCtxt@ very similar to @Level 0 0@, but is used for one purpose:
to say "don't float anything out of here".  That's exactly what we
want for the body of an INLINE, where we don't want to float anything
out at all.  See notes with lvlMFE below.

But, check this out:

-- At one time I tried the effect of not float anything out of an InlineMe,
-- but it sometimes works badly.  For example, consider PrelArr.done.  It
-- has the form 	__inline (\d. e)
-- where e doesn't mention d.  If we float this to 
--	__inline (let x = e in \d. x)
-- things are bad.  The inliner doesn't even inline it because it doesn't look
-- like a head-normal form.  So it seems a lesser evil to let things float.
-- In SetLevels we do set the context to (Level 0 0) when we get to an InlineMe
-- which discourages floating out.

So the conclusion is: don't do any floating at all inside an InlineMe.
(In the above example, don't float the {x=e} out of the \d.)

One particular case is that of workers: we don't want to float the
call to the worker outside the wrapper, otherwise the worker might get
inlined into the floated expression, and an importing module won't see
the worker at all.

149 150 151 152
\begin{code}
type LevelledExpr  = TaggedExpr Level
type LevelledBind  = TaggedBind Level

153 154
tOP_LEVEL   = Level 0 0
iNLINE_CTXT = InlineCtxt
155 156

incMajorLvl :: Level -> Level
157
-- For InlineCtxt we ignore any inc's; we don't want
158
-- to do any floating at all; see notes above
159
incMajorLvl InlineCtxt		= InlineCtxt
160 161 162
incMajorLvl (Level major minor) = Level (major+1) 0

incMinorLvl :: Level -> Level
163
incMinorLvl InlineCtxt		= InlineCtxt
164 165 166
incMinorLvl (Level major minor) = Level major (minor+1)

maxLvl :: Level -> Level -> Level
167 168
maxLvl InlineCtxt l2  = l2
maxLvl l1  InlineCtxt = l1
169 170 171 172 173
maxLvl l1@(Level maj1 min1) l2@(Level maj2 min2)
  | (maj1 > maj2) || (maj1 == maj2 && min1 > min2) = l1
  | otherwise					   = l2

ltLvl :: Level -> Level -> Bool
174 175
ltLvl any_lvl	 InlineCtxt  = False
ltLvl InlineCtxt (Level _ _) = True
176 177 178 179 180
ltLvl (Level maj1 min1) (Level maj2 min2)
  = (maj1 < maj2) || (maj1 == maj2 && min1 < min2)

ltMajLvl :: Level -> Level -> Bool
    -- Tells if one level belongs to a difft *lambda* level to another
181 182
ltMajLvl any_lvl	InlineCtxt     = False
ltMajLvl InlineCtxt	(Level maj2 _) = 0 < maj2
183 184 185 186
ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2

isTopLvl :: Level -> Bool
isTopLvl (Level 0 0) = True
187 188 189 190 191
isTopLvl other	     = False

isInlineCtxt :: Level -> Bool
isInlineCtxt InlineCtxt = True
isInlineCtxt other	= False
192 193

instance Outputable Level where
194
  ppr InlineCtxt      = text "<INLINE>"
195 196 197
  ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]

instance Eq Level where
198
  InlineCtxt	    == InlineCtxt	 = True
199
  (Level maj1 min1) == (Level maj2 min2) = maj1==maj2 && min1==min2
200
  l1		    == l2		 = False
201 202
\end{code}

203

204 205 206 207 208 209 210
%************************************************************************
%*									*
\subsection{Main level-setting code}
%*									*
%************************************************************************

\begin{code}
211
setLevels :: FloatOutSwitches
212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277
	  -> [CoreBind]
	  -> UniqSupply
	  -> [LevelledBind]

setLevels float_lams binds us
  = initLvl us (do_them binds)
  where
    -- "do_them"'s main business is to thread the monad along
    -- It gives each top binding the same empty envt, because
    -- things unbound in the envt have level number zero implicitly
    do_them :: [CoreBind] -> LvlM [LevelledBind]

    do_them [] = returnLvl []
    do_them (b:bs)
      = lvlTopBind init_env b	`thenLvl` \ (lvld_bind, _) ->
	do_them bs		`thenLvl` \ lvld_binds ->
    	returnLvl (lvld_bind : lvld_binds)

    init_env = initialEnv float_lams

lvlTopBind env (NonRec binder rhs)
  = lvlBind TopLevel tOP_LEVEL env (AnnNonRec binder (freeVars rhs))
					-- Rhs can have no free vars!

lvlTopBind env (Rec pairs)
  = lvlBind TopLevel tOP_LEVEL env (AnnRec [(b,freeVars rhs) | (b,rhs) <- pairs])
\end{code}

%************************************************************************
%*									*
\subsection{Setting expression levels}
%*									*
%************************************************************************

\begin{code}
lvlExpr :: Level		-- ctxt_lvl: Level of enclosing expression
	-> LevelEnv		-- Level of in-scope names/tyvars
	-> CoreExprWithFVs	-- input expression
	-> LvlM LevelledExpr	-- Result expression
\end{code}

The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
binder.  Here's an example

	v = \x -> ...\y -> let r = case (..x..) of
					..x..
			   in ..

When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
the level of @r@, even though it's inside a level-2 @\y@.  It's
important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
--- because it isn't a *maximal* free expression.

If there were another lambda in @r@'s rhs, it would get level-2 as well.

\begin{code}
lvlExpr _ _ (_, AnnType ty)   = returnLvl (Type ty)
lvlExpr _ env (_, AnnVar v)   = returnLvl (lookupVar env v)
lvlExpr _ env (_, AnnLit lit) = returnLvl (Lit lit)

lvlExpr ctxt_lvl env (_, AnnApp fun arg)
  = lvl_fun fun				`thenLvl` \ fun' ->
    lvlMFE  False ctxt_lvl env arg	`thenLvl` \ arg' ->
    returnLvl (App fun' arg')
  where
278 279
-- gaw 2004
    lvl_fun (_, AnnCase _ _ _ _) = lvlMFE True ctxt_lvl env fun
280 281 282 283 284
    lvl_fun other 	       = lvlExpr ctxt_lvl env fun
	-- We don't do MFE on partial applications generally,
	-- but we do if the function is big and hairy, like a case

lvlExpr ctxt_lvl env (_, AnnNote InlineMe expr)
285 286
-- Don't float anything out of an InlineMe; hence the iNLINE_CTXT
  = lvlExpr iNLINE_CTXT env expr 	`thenLvl` \ expr' ->
287 288 289 290 291 292 293 294 295 296 297 298 299 300 301
    returnLvl (Note InlineMe expr')

lvlExpr ctxt_lvl env (_, AnnNote note expr)
  = lvlExpr ctxt_lvl env expr 		`thenLvl` \ expr' ->
    returnLvl (Note note expr')

-- We don't split adjacent lambdas.  That is, given
--	\x y -> (x+1,y)
-- we don't float to give 
--	\x -> let v = x+y in \y -> (v,y)
-- Why not?  Because partial applications are fairly rare, and splitting
-- lambdas makes them more expensive.

lvlExpr ctxt_lvl env expr@(_, AnnLam bndr rhs)
  = lvlMFE True new_lvl new_env body	`thenLvl` \ new_body ->
302
    returnLvl (mkLams new_bndrs new_body)
303
  where 
304
    (bndrs, body)	 = collectAnnBndrs expr
305 306
    (new_lvl, new_bndrs) = lvlLamBndrs ctxt_lvl bndrs
    new_env 		 = extendLvlEnv env new_bndrs
307 308 309 310 311 312
	-- At one time we called a special verion of collectBinders,
	-- which ignored coercions, because we don't want to split
	-- a lambda like this (\x -> coerce t (\s -> ...))
	-- This used to happen quite a bit in state-transformer programs,
	-- but not nearly so much now non-recursive newtypes are transparent.
	-- [See SetLevels rev 1.50 for a version with this approach.]
313

314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330
lvlExpr ctxt_lvl env (_, AnnLet (AnnNonRec bndr rhs) body)
  | isUnLiftedType (idType bndr)
	-- Treat unlifted let-bindings (let x = b in e) just like (case b of x -> e)
	-- That is, leave it exactly where it is
	-- We used to float unlifted bindings too (e.g. to get a cheap primop
	-- outside a lambda (to see how, look at lvlBind in rev 1.58)
	-- but an unrelated change meant that these unlifed bindings
	-- could get to the top level which is bad.  And there's not much point;
	-- unlifted bindings are always cheap, and so hardly worth floating.
  = lvlExpr ctxt_lvl env rhs		`thenLvl` \ rhs' ->
    lvlExpr incd_lvl env' body		`thenLvl` \ body' ->
    returnLvl (Let (NonRec bndr' rhs') body')
  where
    incd_lvl = incMinorLvl ctxt_lvl
    bndr' = TB bndr incd_lvl
    env'  = extendLvlEnv env [bndr']

331 332 333 334 335
lvlExpr ctxt_lvl env (_, AnnLet bind body)
  = lvlBind NotTopLevel ctxt_lvl env bind	`thenLvl` \ (bind', new_env) ->
    lvlExpr ctxt_lvl new_env body		`thenLvl` \ body' ->
    returnLvl (Let bind' body')

336
lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr ty alts)
337 338 339 340 341
  = lvlMFE True ctxt_lvl env expr	`thenLvl` \ expr' ->
    let
	alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
    in
    mapLvl (lvl_alt alts_env) alts	`thenLvl` \ alts' ->
342
    returnLvl (Case expr' (TB case_bndr incd_lvl) ty alts')
343 344 345 346 347 348 349
  where
      incd_lvl  = incMinorLvl ctxt_lvl

      lvl_alt alts_env (con, bs, rhs)
	= lvlMFE True incd_lvl new_env rhs	`thenLvl` \ rhs' ->
	  returnLvl (con, bs', rhs')
	where
350
	  bs'     = [ TB b incd_lvl | b <- bs ]
351 352 353 354 355 356
	  new_env = extendLvlEnv alts_env bs'
\end{code}

@lvlMFE@ is just like @lvlExpr@, except that it might let-bind
the expression, so that it can itself be floated.

357 358 359 360 361 362 363
[NOTE: unlifted MFEs]
We don't float unlifted MFEs, which potentially loses big opportunites.
For example:
	\x -> f (h y)
where h :: Int -> Int# is expensive. We'd like to float the (h y) outside
the \x, but we don't because it's unboxed.  Possible solution: box it.

364 365 366 367 368 369 370 371 372 373
\begin{code}
lvlMFE ::  Bool			-- True <=> strict context [body of case or let]
	-> Level		-- Level of innermost enclosing lambda/tylam
	-> LevelEnv		-- Level of in-scope names/tyvars
	-> CoreExprWithFVs	-- input expression
	-> LvlM LevelledExpr	-- Result expression

lvlMFE strict_ctxt ctxt_lvl env (_, AnnType ty)
  = returnLvl (Type ty)

374

375
lvlMFE strict_ctxt ctxt_lvl env ann_expr@(fvs, _)
376
  |  isUnLiftedType ty			-- Can't let-bind it; see [NOTE: unlifted MFEs]
377
  || isInlineCtxt ctxt_lvl		-- Don't float out of an __inline__ context
378
  || exprIsTrivial expr			-- Never float if it's trivial
379 380 381 382 383 384 385
  || not good_destination
  = 	-- Don't float it out
    lvlExpr ctxt_lvl env ann_expr

  | otherwise	-- Float it out!
  = lvlFloatRhs abs_vars dest_lvl env ann_expr	`thenLvl` \ expr' ->
    newLvlVar "lvl" abs_vars ty			`thenLvl` \ var ->
386
    returnLvl (Let (NonRec (TB var dest_lvl) expr') 
387 388 389 390 391 392 393 394 395
		   (mkVarApps (Var var) abs_vars))
  where
    expr     = deAnnotate ann_expr
    ty       = exprType expr
    dest_lvl = destLevel env fvs (isFunction ann_expr)
    abs_vars = abstractVars dest_lvl env fvs

	-- A decision to float entails let-binding this thing, and we only do 
	-- that if we'll escape a value lambda, or will go to the top level.
396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429
    good_destination 
	| dest_lvl `ltMajLvl` ctxt_lvl		-- Escapes a value lambda
	= not (exprIsCheap expr) || isTopLvl dest_lvl
	  -- Even if it escapes a value lambda, we only
	  -- float if it's not cheap (unless it'll get all the
	  -- way to the top).  I've seen cases where we
	  -- float dozens of tiny free expressions, which cost
	  -- more to allocate than to evaluate.
	  -- NB: exprIsCheap is also true of bottom expressions, which
	  --     is good; we don't want to share them
	  --
	  -- It's only Really Bad to float a cheap expression out of a
	  -- strict context, because that builds a thunk that otherwise
	  -- would never be built.  So another alternative would be to
	  -- add 
	  -- 	|| (strict_ctxt && not (exprIsBottom expr))
	  -- to the condition above. We should really try this out.

	| otherwise		-- Does not escape a value lambda
	= isTopLvl dest_lvl 	-- Only float if we are going to the top level
	&& floatConsts env	--   and the floatConsts flag is on
	&& not strict_ctxt	-- Don't float from a strict context	
	  -- We are keen to float something to the top level, even if it does not
	  -- escape a lambda, because then it needs no allocation.  But it's controlled
	  -- by a flag, because doing this too early loses opportunities for RULES
	  -- which (needless to say) are important in some nofib programs
	  -- (gcd is an example).
	  --
	  -- Beware:
	  --	concat = /\ a -> foldr ..a.. (++) []
	  -- was getting turned into
	  --	concat = /\ a -> lvl a
	  --	lvl    = /\ a -> foldr ..a.. (++) []
	  -- which is pretty stupid.  Hence the strict_ctxt test
430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449
\end{code}


%************************************************************************
%*									*
\subsection{Bindings}
%*									*
%************************************************************************

The binding stuff works for top level too.

\begin{code}
lvlBind :: TopLevelFlag		-- Used solely to decide whether to clone
	-> Level		-- Context level; might be Top even for bindings nested in the RHS
				-- of a top level binding
	-> LevelEnv
	-> CoreBindWithFVs
	-> LvlM (LevelledBind, LevelEnv)

lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
450
  | isInlineCtxt ctxt_lvl		-- Don't do anything inside InlineMe
451
  = lvlExpr ctxt_lvl env rhs			`thenLvl` \ rhs' ->
452
    returnLvl (NonRec (TB bndr ctxt_lvl) rhs', env)
453

454 455 456 457
  | null abs_vars
  =	-- No type abstraction; clone existing binder
    lvlExpr dest_lvl env rhs			`thenLvl` \ rhs' ->
    cloneVar top_lvl env bndr ctxt_lvl dest_lvl	`thenLvl` \ (env', bndr') ->
458
    returnLvl (NonRec (TB bndr' dest_lvl) rhs', env') 
459 460 461 462 463

  | otherwise
  = -- Yes, type abstraction; create a new binder, extend substitution, etc
    lvlFloatRhs abs_vars dest_lvl env rhs	`thenLvl` \ rhs' ->
    newPolyBndrs dest_lvl env abs_vars [bndr]	`thenLvl` \ (env', [bndr']) ->
464
    returnLvl (NonRec (TB bndr' dest_lvl) rhs', env')
465 466 467 468

  where
    bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
    abs_vars = abstractVars dest_lvl env bind_fvs
469
    dest_lvl = destLevel env bind_fvs (isFunction rhs)
470 471 472 473 474
\end{code}


\begin{code}
lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
475 476
  | isInlineCtxt ctxt_lvl	-- Don't do anything inside InlineMe
  = mapLvl (lvlExpr ctxt_lvl env) rhss			`thenLvl` \ rhss' ->
477
    returnLvl (Rec ([TB b ctxt_lvl | b <- bndrs] `zip` rhss'), env)
478

479 480 481
  | null abs_vars
  = cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl	`thenLvl` \ (new_env, new_bndrs) ->
    mapLvl (lvlExpr ctxt_lvl new_env) rhss		`thenLvl` \ new_rhss ->
482
    returnLvl (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501

  | isSingleton pairs && count isId abs_vars > 1
  = 	-- Special case for self recursion where there are
	-- several variables carried around: build a local loop:	
	--	poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
	-- This just makes the closures a bit smaller.  If we don't do
	-- this, allocation rises significantly on some programs
	--
	-- We could elaborate it for the case where there are several
	-- mutually functions, but it's quite a bit more complicated
	-- 
	-- This all seems a bit ad hoc -- sigh
    let
	(bndr,rhs) = head pairs
	(rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
	rhs_env = extendLvlEnv env abs_vars_w_lvls
    in
    cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl	`thenLvl` \ (rhs_env', new_bndr) ->
    let
502
	(lam_bndrs, rhs_body)     = collectAnnBndrs rhs
503 504 505 506 507
        (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
	body_env 		  = extendLvlEnv rhs_env' new_lam_bndrs
    in
    lvlExpr body_lvl body_env rhs_body		`thenLvl` \ new_rhs_body ->
    newPolyBndrs dest_lvl env abs_vars [bndr]	`thenLvl` \ (poly_env, [poly_bndr]) ->
508 509 510 511 512
    returnLvl (Rec [(TB poly_bndr dest_lvl, 
	       mkLams abs_vars_w_lvls $
	       mkLams new_lam_bndrs $
	       Let (Rec [(TB new_bndr rhs_lvl, mkLams new_lam_bndrs new_rhs_body)]) 
		   (mkVarApps (Var new_bndr) lam_bndrs))],
513 514
	       poly_env)

515
  | otherwise	-- Non-null abs_vars
516 517
  = newPolyBndrs dest_lvl env abs_vars bndrs		`thenLvl` \ (new_env, new_bndrs) ->
    mapLvl (lvlFloatRhs abs_vars dest_lvl new_env) rhss `thenLvl` \ new_rhss ->
518
    returnLvl (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550

  where
    (bndrs,rhss) = unzip pairs

	-- Finding the free vars of the binding group is annoying
    bind_fvs	    = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
				    | (bndr, (rhs_fvs,_)) <- pairs])
		      `minusVarSet`
		      mkVarSet bndrs

    dest_lvl = destLevel env bind_fvs (all isFunction rhss)
    abs_vars = abstractVars dest_lvl env bind_fvs

----------------------------------------------------
-- Three help functons for the type-abstraction case

lvlFloatRhs abs_vars dest_lvl env rhs
  = lvlExpr rhs_lvl rhs_env rhs	`thenLvl` \ rhs' ->
    returnLvl (mkLams abs_vars_w_lvls rhs')
  where
    (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
    rhs_env = extendLvlEnv env abs_vars_w_lvls
\end{code}


%************************************************************************
%*									*
\subsection{Deciding floatability}
%*									*
%************************************************************************

\begin{code}
551
lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [TaggedBndr Level])
552 553 554 555 556 557 558 559 560 561 562 563 564 565 566
-- Compute the levels for the binders of a lambda group
-- The binders returned are exactly the same as the ones passed,
-- but they are now paired with a level
lvlLamBndrs lvl [] 
  = (lvl, [])

lvlLamBndrs lvl bndrs
  = go  (incMinorLvl lvl)
	False 	-- Havn't bumped major level in this group
	[] bndrs
  where
    go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
	| isId bndr && 			-- Go to the next major level if this is a value binder,
	  not bumped_major && 		-- and we havn't already gone to the next level (one jump per group)
	  not (isOneShotLambda bndr)	-- and it isn't a one-shot lambda
567
	= go new_lvl True (TB bndr new_lvl : rev_lvld_bndrs) bndrs
568 569

	| otherwise
570
	= go old_lvl bumped_major (TB bndr old_lvl : rev_lvld_bndrs) bndrs
571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619

	where
	  new_lvl = incMajorLvl old_lvl

    go old_lvl _ rev_lvld_bndrs []
	= (old_lvl, reverse rev_lvld_bndrs)
	-- a lambda like this (\x -> coerce t (\s -> ...))
	-- This happens quite a bit in state-transformer programs
\end{code}

\begin{code}
  -- Destintion level is the max Id level of the expression
  -- (We'll abstract the type variables, if any.)
destLevel :: LevelEnv -> VarSet -> Bool -> Level
destLevel env fvs is_function
  |  floatLams env
  && is_function = tOP_LEVEL		-- Send functions to top level; see
					-- the comments with isFunction
  | otherwise    = maxIdLevel env fvs

isFunction :: CoreExprWithFVs -> Bool
-- The idea here is that we want to float *functions* to
-- the top level.  This saves no work, but 
--	(a) it can make the host function body a lot smaller, 
--		and hence inlinable.  
--	(b) it can also save allocation when the function is recursive:
--	    h = \x -> letrec f = \y -> ...f...y...x...
--		      in f x
--     becomes
--	    f = \x y -> ...(f x)...y...x...
--	    h = \x -> f x x
--     No allocation for f now.
-- We may only want to do this if there are sufficiently few free 
-- variables.  We certainly only want to do it for values, and not for
-- constructors.  So the simple thing is just to look for lambdas
isFunction (_, AnnLam b e) | isId b    = True
			   | otherwise = isFunction e
isFunction (_, AnnNote n e)            = isFunction e
isFunction other 		       = False
\end{code}


%************************************************************************
%*									*
\subsection{Free-To-Level Monad}
%*									*
%************************************************************************

\begin{code}
620
type LevelEnv = (FloatOutSwitches,
621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645
		 VarEnv Level, 			-- Domain is *post-cloned* TyVars and Ids
	         Subst, 			-- Domain is pre-cloned Ids; tracks the in-scope set
						-- 	so that subtitution is capture-avoiding
	         IdEnv ([Var], LevelledExpr))	-- Domain is pre-cloned Ids
	-- We clone let-bound variables so that they are still
	-- distinct when floated out; hence the SubstEnv/IdEnv.
        -- (see point 3 of the module overview comment).
	-- We also use these envs when making a variable polymorphic
	-- because we want to float it out past a big lambda.
	--
	-- The SubstEnv and IdEnv always implement the same mapping, but the
	-- SubstEnv maps to CoreExpr and the IdEnv to LevelledExpr
	-- Since the range is always a variable or type application,
	-- there is never any difference between the two, but sadly
	-- the types differ.  The SubstEnv is used when substituting in
	-- a variable's IdInfo; the IdEnv when we find a Var.
	--
	-- In addition the IdEnv records a list of tyvars free in the
	-- type application, just so we don't have to call freeVars on
	-- the type application repeatedly.
	--
	-- The domain of the both envs is *pre-cloned* Ids, though
	--
	-- The domain of the VarEnv Level is the *post-cloned* Ids

646
initialEnv :: FloatOutSwitches -> LevelEnv
647 648 649
initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)

floatLams :: LevelEnv -> Bool
650 651 652 653
floatLams (FloatOutSw float_lams _, _, _, _) = float_lams

floatConsts :: LevelEnv -> Bool
floatConsts (FloatOutSw _ float_consts, _, _, _) = float_consts
654

655
extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
656 657 658 659 660 661 662
-- Used when *not* cloning
extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
  = (float_lams,
     foldl add_lvl lvl_env prs,
     foldl del_subst subst prs,
     foldl del_id id_env prs)
  where
663 664 665
    add_lvl   env (TB v l) = extendVarEnv env v l
    del_subst env (TB v _) = extendInScope env v
    del_id    env (TB v _) = delVarEnv env v
666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684
  -- We must remove any clone for this variable name in case of
  -- shadowing.  This bit me in the following case
  -- (in nofib/real/gg/Spark.hs):
  -- 
  --   case ds of wild {
  --     ... -> case e of wild {
  --              ... -> ... wild ...
  --            }
  --   }
  -- 
  -- The inside occurrence of @wild@ was being replaced with @ds@,
  -- incorrectly, because the SubstEnv was still lying around.  Ouch!
  -- KSW 2000-07.

-- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
-- (see point 4 of the module overview comment)
extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
  = (float_lams,
     extendVarEnv lvl_env case_bndr lvl,
685
     extendIdSubst subst case_bndr (Var scrut_var),
686 687 688
     extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
     
extendCaseBndrLvlEnv env scrut case_bndr lvl
689
  = extendLvlEnv          env [TB case_bndr lvl]
690 691 692 693 694 695 696 697

extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
  = (float_lams,
     foldl add_lvl   lvl_env bndr_pairs,
     foldl add_subst subst   bndr_pairs,
     foldl add_id    id_env  bndr_pairs)
  where
     add_lvl   env (v,v') = extendVarEnv env v' dest_lvl
698
     add_subst env (v,v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729
     add_id    env (v,v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)

extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
  = (float_lams,
     foldl add_lvl   lvl_env bndr_pairs,
     new_subst,
     foldl add_id    id_env  bndr_pairs)
  where
     add_lvl   env (v,v') = extendVarEnv env v' lvl
     add_id    env (v,v') = extendVarEnv env v ([v'], Var v')


maxIdLevel :: LevelEnv -> VarSet -> Level
maxIdLevel (_, lvl_env,_,id_env) var_set
  = foldVarSet max_in tOP_LEVEL var_set
  where
    max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
						Just (abs_vars, _) -> abs_vars
						Nothing		   -> [in_var])

    max_out out_var lvl 
	| isId out_var = case lookupVarEnv lvl_env out_var of
				Just lvl' -> maxLvl lvl' lvl
				Nothing   -> lvl 
	| otherwise    = lvl	-- Ignore tyvars in *maxIdLevel*

lookupVar :: LevelEnv -> Id -> LevelledExpr
lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
				       Just (_, expr) -> expr
				       other	      -> Var v

730 731 732 733 734
abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
	-- Find the variables in fvs, free vars of the target expresion,
	-- whose level is greater than the destination level
	-- These are the ones we are going to abstract out
abstractVars dest_lvl env fvs
735
  = uniq (sortLe le [var | fv <- varSetElems fvs, var <- absVarsOf dest_lvl env fv])
736 737 738
  where
	-- Sort the variables so we don't get 
	-- mixed-up tyvars and Ids; it's just messy
739
    v1 `le` v2 = case (isId v1, isId v2) of
740 741
		   (True, False) -> False
		   (False, True) -> True
742
		   other	 -> v1 <= v2	-- Same family
743 744 745 746 747 748 749

    uniq :: [Var] -> [Var]
	-- Remove adjacent duplicates; the sort will have brought them together
    uniq (v1:v2:vs) | v1 == v2  = uniq (v2:vs)
		    | otherwise = v1 : uniq (v2:vs)
    uniq vs = vs

750
absVarsOf :: Level -> LevelEnv -> Var -> [Var]
751
	-- If f is free in the expression, and f maps to poly_f a b c in the
752 753 754 755
	-- current substitution, then we must report a b c as candidate type
	-- variables
absVarsOf dest_lvl (_, lvl_env, _, id_env) v 
  | isId v
756
  = [zap av2 | av1 <- lookup_avs v, av2 <- add_tyvars av1, abstract_me av2]
757 758 759 760 761 762 763 764 765 766 767 768 769

  | otherwise
  = if abstract_me v then [v] else []

  where
    abstract_me v = case lookupVarEnv lvl_env v of
			Just lvl -> dest_lvl `ltLvl` lvl
			Nothing  -> False

    lookup_avs v = case lookupVarEnv id_env v of
			Just (abs_vars, _) -> abs_vars
			Nothing	           -> [v]

770
    add_tyvars v | isId v    = v : varSetElems (idFreeTyVars v)
771 772
		 | otherwise = [v]

773 774 775 776 777 778 779
	-- We are going to lambda-abstract, so nuke any IdInfo,
	-- and add the tyvars of the Id (if necessary)
    zap v | isId v = WARN( workerExists (idWorkerInfo v) ||
		           not (isEmptyCoreRules (idSpecialisation v)),
		           text "absVarsOf: discarding info on" <+> ppr v )
		     setIdInfo v vanillaIdInfo
	  | otherwise = v
780 781 782 783 784 785 786 787 788 789 790 791 792
\end{code}

\begin{code}
type LvlM result = UniqSM result

initLvl		= initUs_
thenLvl		= thenUs
returnLvl	= returnUs
mapLvl		= mapUs
\end{code}

\begin{code}
newPolyBndrs dest_lvl env abs_vars bndrs
793
  = getUniquesUs 		`thenLvl` \ uniqs ->
794 795 796 797 798
    let
	new_bndrs = zipWith mk_poly_bndr bndrs uniqs
    in
    returnLvl (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
  where
799
    mk_poly_bndr bndr uniq = mkSysLocalUnencoded (mkFastString str) uniq poly_ty
800 801
			   where
			     str     = "poly_" ++ occNameUserString (getOccName bndr)
802
			     poly_ty = mkPiTypes abs_vars (idType bndr)
803 804 805 806 807 808 809
	

newLvlVar :: String 
	  -> [CoreBndr] -> Type 	-- Abstract wrt these bndrs
	  -> LvlM Id
newLvlVar str vars body_ty 	
  = getUniqueUs	`thenLvl` \ uniq ->
810
    returnUs (mkSysLocalUnencoded (mkFastString str) uniq (mkPiTypes vars body_ty))
811 812 813 814 815 816 817 818 819 820 821
    
-- The deeply tiresome thing is that we have to apply the substitution
-- to the rules inside each Id.  Grr.  But it matters.

cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
cloneVar TopLevel env v ctxt_lvl dest_lvl
  = returnUs (env, v)	-- Don't clone top level things
cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
  = ASSERT( isId v )
    getUs	`thenLvl` \ us ->
    let
822
      (subst', v1) = cloneIdBndr subst us v
823 824 825 826 827 828 829 830 831 832 833 834
      v2	   = zap_demand ctxt_lvl dest_lvl v1
      env'	   = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
    in
    returnUs (env', v2)

cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
cloneRecVars TopLevel env vs ctxt_lvl dest_lvl 
  = returnUs (env, vs)	-- Don't clone top level things
cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
  = ASSERT( all isId vs )
    getUs 			`thenLvl` \ us ->
    let
835
      (subst', vs1) = cloneRecIdBndrs subst us vs
836 837 838 839 840 841 842 843 844 845 846 847
      vs2	    = map (zap_demand ctxt_lvl dest_lvl) vs1
      env'	    = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
    in
    returnUs (env', vs2)

	-- VERY IMPORTANT: we must zap the demand info 
	-- if the thing is going to float out past a lambda
zap_demand dest_lvl ctxt_lvl id
  | ctxt_lvl == dest_lvl = id			-- Stays put
  | otherwise		 = zapDemandIdInfo id	-- Floats out
\end{code}