LiberateCase.lhs 12.5 KB
Newer Older
1
%
2
% (c) The AQUA Project, Glasgow University, 1994-1998
3 4 5 6
%
\section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}

\begin{code}
Ian Lynagh's avatar
Ian Lynagh committed
7 8 9 10 11 12 13
{-# OPTIONS -fno-warn-tabs #-}
-- The above warning supression flag is a temporary kludge.
-- While working on this module you are encouraged to remove it and
-- detab the module (please do the detabbing in a separate patch). See
--     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#TabsvsSpaces
-- for details

14 15
module LiberateCase ( liberateCase ) where

16 17
#include "HsVersions.h"

18
import DynFlags
19
import CoreSyn
20
import CoreUnfold	( couldBeSmallEnoughToInline )
21
import Id
22
import VarEnv
sof's avatar
sof committed
23
import Util             ( notNull )
24 25
\end{code}

26 27
The liberate-case transformation
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
28 29 30 31 32 33 34
This module walks over @Core@, and looks for @case@ on free variables.
The criterion is:
	if there is case on a free on the route to the recursive call,
	then the recursive call is replaced with an unfolding.

Example

35 36
   f = \ t -> case v of
	         V a b -> a : f t
37 38 39

=> the inner f is replaced.

40 41
   f = \ t -> case v of
	         V a b -> a : (letrec
42
				f =  \ t -> case v of
43
					       V a b -> a : f t
44
			       in f) t
45 46
(note the NEED for shadowing)

47 48
=> Simplify

49 50
  f = \ t -> case v of
	         V a b -> a : (letrec
51
				f = \ t -> a : f t
52
			       in f t)
53

54 55 56
Better code, because 'a' is  free inside the inner letrec, rather
than needing projection from v.

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
57 58
Note that this deals with *free variables*.  SpecConstr deals with
*arguments* that are of known form.  E.g.
59 60 61 62 63 64

	last []     = error 
	last (x:[]) = x
	last (x:xs) = last xs

	
65 66 67 68 69 70
Note [Scrutinee with cast]
~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this:
    f = \ t -> case (v `cast` co) of
	         V a b -> a : f t

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
71
Exactly the same optimisation (unrolling one call to f) will work here, 
72 73
despite the cast.  See mk_alt_env in the Case branch of libCase.

74

75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92
Note [Only functions!]
~~~~~~~~~~~~~~~~~~~~~~
Consider the following code

       f = g (case v of V a b -> a : t f)

where g is expensive. If we aren't careful, liberate case will turn this into

       f = g (case v of
               V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
                                in f)
             )

Yikes! We evaluate g twice. This leads to a O(2^n) explosion
if g calls back to the same code recursively.

Solution: make sure that we only do the liberate-case thing on *functions*

93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115
To think about (Apr 94)
~~~~~~~~~~~~~~
Main worry: duplicating code excessively.  At the moment we duplicate
the entire binding group once at each recursive call.  But there may
be a group of recursive calls which share a common set of evaluated
free variables, in which case the duplication is a plain waste.

Another thing we could consider adding is some unfold-threshold thing,
so that we'll only duplicate if the size of the group rhss isn't too
big.

Data types
~~~~~~~~~~
The ``level'' of a binder tells how many
recursive defns lexically enclose the binding
A recursive defn "encloses" its RHS, not its
scope.  For example:
\begin{verbatim}
	letrec f = let g = ... in ...
	in
	let h = ...
	in ...
\end{verbatim}
116
Here, the level of @f@ is zero, the level of @g@ is one,
117 118
and the level of @h@ is zero (NB not one).

119 120 121 122 123 124

%************************************************************************
%*									*
	 Top-level code
%*									*
%************************************************************************
125 126

\begin{code}
127
liberateCase :: DynFlags -> CoreProgram -> CoreProgram
128
liberateCase dflags binds = do_prog (initEnv dflags) binds
129
  where
Ian Lynagh's avatar
Ian Lynagh committed
130
    do_prog _   [] = []
131 132 133 134 135
    do_prog env (bind:binds) = bind' : do_prog env' binds
			     where
			       (env', bind') = libCaseBind env bind
\end{code}

136 137 138 139 140 141 142

%************************************************************************
%*									*
	 Main payload
%*									*
%************************************************************************

143 144 145
Bindings
~~~~~~~~
\begin{code}
146
libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
147

148 149
libCaseBind env (NonRec binder rhs)
  = (addBinders env [binder], NonRec binder (libCase env rhs))
150

151 152
libCaseBind env (Rec pairs)
  = (env_body, Rec pairs')
153
  where
154
    binders = map fst pairs
155 156 157 158 159 160 161 162

    env_body = addBinders env binders

    pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]

	-- We extend the rec-env by binding each Id to its rhs, first
	-- processing the rhs with an *un-extended* environment, so
	-- that the same process doesn't occur for ever!
163 164 165 166 167 168 169
    env_rhs = addRecBinds env [ (localiseId binder, libCase env_body rhs)
			      | (binder, rhs) <- pairs
 			      , rhs_small_enough binder rhs ]
	-- localiseID : see Note [Need to localiseId in libCaseBind]
		 

    rhs_small_enough id rhs	-- Note [Small enough]
170
	=  idArity id > 0	-- Note [Only functions!]
171
	&& maybe True (\size -> couldBeSmallEnoughToInline (lc_dflags env) size rhs)
172
                      (bombOutSize env)
173 174
\end{code}

175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196
Note [Need to localiseId in libCaseBind]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The call to localiseId is needed for two subtle reasons
(a)  Reset the export flags on the binders so
	that we don't get name clashes on exported things if the 
	local binding floats out to top level.  This is most unlikely
	to happen, since the whole point concerns free variables. 
	But resetting the export flag is right regardless.

(b)  Make the name an Internal one.  External Names should never be
	nested; if it were floated to the top level, we'd get a name
	clash at code generation time.

Note [Small enough]
~~~~~~~~~~~~~~~~~~~
Consider
  \fv. letrec
     	 f = \x. BIG...(case fv of { (a,b) -> ...g.. })...
     	 g = \y. SMALL...f...
Then we *can* do liberate-case on g (small RHS) but not for f (too big).
But we can choose on a item-by-item basis, and that's what the
rhs_small_enough call in the comprehension for env_rhs does.
197 198 199 200 201 202

Expressions
~~~~~~~~~~~

\begin{code}
libCase :: LibCaseEnv
203 204
	-> CoreExpr
	-> CoreExpr
205

Ian Lynagh's avatar
Ian Lynagh committed
206 207 208
libCase env (Var v)             = libCaseId env v
libCase _   (Lit lit)           = Lit lit
libCase _   (Type ty)           = Type ty
209
libCase _   (Coercion co)       = Coercion co
210
libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
211
libCase env (Tick tickish body) = Tick tickish (libCase env body)
212
libCase env (Cast e co)         = Cast (libCase env e) co
213

214 215
libCase env (Lam binder body)
  = Lam binder (libCase (addBinders env [binder]) body)
216

217 218
libCase env (Let bind body)
  = Let bind' (libCase env_body body)
219 220 221
  where
    (env_body, bind') = libCaseBind env bind

222
libCase env (Case scrut bndr ty alts)
223
  = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
224
  where
225 226 227
    env_alts = addBinders (mk_alt_env scrut) [bndr]
    mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
    mk_alt_env (Cast scrut _)  = mk_alt_env scrut	-- Note [Scrutinee with cast]
Ian Lynagh's avatar
Ian Lynagh committed
228
    mk_alt_env _    	       = env
229

Ian Lynagh's avatar
Ian Lynagh committed
230 231
libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
                         -> (AltCon, [CoreBndr], CoreExpr)
232
libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
233 234
\end{code}

235

236 237
Ids
~~~
238
\begin{code}
239
libCaseId :: LibCaseEnv -> Id -> CoreExpr
240
libCaseId env v
241
  | Just the_bind <- lookupRecId env v	-- It's a use of a recursive thing
sof's avatar
sof committed
242
  , notNull free_scruts 		-- with free vars scrutinised in RHS
243
  = Let the_bind (Var v)
244 245

  | otherwise
246
  = Var v
247 248

  where
249 250
    rec_id_level = lookupLevel env v
    free_scruts  = freeScruts env rec_id_level
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
251 252 253 254 255 256

freeScruts :: LibCaseEnv
	   -> LibCaseLevel 	-- Level of the recursive Id
	   -> [Id]		-- Ids that are scrutinised between the binding
				-- of the recursive Id and here
freeScruts env rec_bind_lvl
257 258 259
  = [v | (v, scrut_bind_lvl, scrut_at_lvl) <- lc_scruts env
       , scrut_bind_lvl <= rec_bind_lvl
       , scrut_at_lvl > rec_bind_lvl]
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
260
	-- Note [When to specialise]
261
	-- Note [Avoiding fruitless liberate-case]
262
\end{code}
263

simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279
Note [When to specialise]
~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  f = \x. letrec g = \y. case x of
			   True  -> ... (f a) ...
			   False -> ... (g b) ...

We get the following levels
	  f  0
	  x  1
	  g  1
	  y  2  

Then 'x' is being scrutinised at a deeper level than its binding, so
it's added to lc_sruts:  [(x,1)]  

Gabor Greif's avatar
typos  
Gabor Greif committed
280
We do *not* want to specialise the call to 'f', because 'x' is not free 
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
281 282 283 284 285
in 'f'.  So here the bind-level of 'x' (=1) is not <= the bind-level of 'f' (=0).

We *do* want to specialise the call to 'g', because 'x' is free in g.
Here the bind-level of 'x' (=1) is <= the bind-level of 'g' (=1).

286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301
Note [Avoiding fruitless liberate-case]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider also:
  f = \x. case top_lvl_thing of
                I# _ -> let g = \y. ... g ...
                        in ...

Here, top_lvl_thing is scrutinised at a level (1) deeper than its
binding site (0).  Nevertheless, we do NOT want to specialise the call
to 'g' because all the structure in its free variables is already
visible at the definition site for g.  Hence, when considering specialising
an occurrence of 'g', we want to check that there's a scruted-var v st

   a) v's binding site is *outside* g
   b) v's scrutinisation site is *inside* g

302

303 304 305 306 307
%************************************************************************
%*									*
	Utility functions
%*									*
%************************************************************************
308 309

\begin{code}
310
addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
311 312
addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
  = env { lc_lvl_env = lvl_env' }
313
  where
314
    lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
315

316
addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
317 318 319
addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env, 
			     lc_rec_env = rec_env}) pairs
  = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
320 321
  where
    lvl'     = lvl + 1
322 323
    lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
    rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
324

325
addScrutedVar :: LibCaseEnv
326
	      -> Id		-- This Id is being scrutinised by a case expression
327
	      -> LibCaseEnv
328

329 330
addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env, 
				lc_scruts = scruts }) scrut_var
331
  | bind_lvl < lvl
332
  = env { lc_scruts = scruts' }
333
	-- Add to scruts iff the scrut_var is being scrutinised at
334
	-- a deeper level than its defn
335 336 337

  | otherwise = env
  where
338
    scruts'  = (scrut_var, bind_lvl, lvl) : scruts
339
    bind_lvl = case lookupVarEnv lvl_env scrut_var of
340
		 Just lvl -> lvl
341
		 Nothing  -> topLevel
342

343
lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
344
lookupRecId env id = lookupVarEnv (lc_rec_env env) id
345 346

lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
347 348
lookupLevel env id
  = case lookupVarEnv (lc_lvl_env env) id of
349
      Just lvl -> lvl
350
      Nothing  -> topLevel
351
\end{code}
352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368

%************************************************************************
%*									*
	 The environment
%*									*
%************************************************************************

\begin{code}
type LibCaseLevel = Int

topLevel :: LibCaseLevel
topLevel = 0
\end{code}

\begin{code}
data LibCaseEnv
  = LibCaseEnv {
369
        lc_dflags :: DynFlags,
370 371

	lc_lvl :: LibCaseLevel,	-- Current level
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
372 373 374
		-- The level is incremented when (and only when) going
		-- inside the RHS of a (sufficiently small) recursive
		-- function.
375 376

	lc_lvl_env :: IdEnv LibCaseLevel,  
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
377 378
		-- Binds all non-top-level in-scope Ids (top-level and
		-- imported things have a level of zero)
379 380

	lc_rec_env :: IdEnv CoreBind, 
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
381 382
		-- Binds *only* recursively defined ids, to their own
		-- binding group, and *only* in their own RHSs
383

384
	lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
385 386
     		-- Each of these Ids was scrutinised by an enclosing
		-- case expression, at a level deeper than its binding
387 388 389 390 391 392 393 394
		-- level.
		-- 
 		-- The first LibCaseLevel is the *binding level* of
 		--   the scrutinised Id, 
		-- The second is the level *at which it was scrutinised*.
		--   (see Note [Avoiding fruitless liberate-case])
		-- The former is a bit redundant, since you could always
		-- look it up in lc_lvl_env, but it's just cached here
simonpj@microsoft.com's avatar
simonpj@microsoft.com committed
395 396
		-- 
		-- The order is insignificant; it's a bag really
397 398 399 400 401
		-- 
		-- There's one element per scrutinisation;
		--    in principle the same Id may appear multiple times,
		--    although that'd be unusual:
		--       case x of { (a,b) -> ....(case x of ...) .. }
402 403
	}

404 405
initEnv :: DynFlags -> LibCaseEnv
initEnv dflags 
406
  = LibCaseEnv { lc_dflags = dflags,
407 408 409
		 lc_lvl = 0,
		 lc_lvl_env = emptyVarEnv, 
		 lc_rec_env = emptyVarEnv,
410
		 lc_scruts = [] }
411

412 413 414
-- Bomb-out size for deciding if
-- potential liberatees are too big.
-- (passed in from cmd-line args)
Ian Lynagh's avatar
Ian Lynagh committed
415
bombOutSize :: LibCaseEnv -> Maybe Int
416
bombOutSize = liberateCaseThreshold . lc_dflags
417 418
\end{code}