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

\begin{code}
module LiberateCase ( liberateCase ) where

9
10
#include "HsVersions.h"

11
import CmdLineOpts	( DynFlags, DynFlag(..), opt_LiberateCaseThreshold )
12
import CoreLint		( showPass, endPass )
13
import CoreSyn
14
import CoreUnfold	( couldBeSmallEnoughToInline )
15
import Id		( Id, setIdName, idName, setIdNotExported )
16
import VarEnv
17
import Name		( localiseName )
18
import Outputable
sof's avatar
sof committed
19
import Util             ( notNull )
20
21
22
23
24
25
26
27
28
29
30
\end{code}

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

\begin{verbatim}
f = \ t -> case v of
31
	       V a b -> a : f t
32
33
34
35
36
37
38
39
\end{verbatim}

=> the inner f is replaced.

\begin{verbatim}
f = \ t -> case v of
	       V a b -> a : (letrec
				f =  \ t -> case v of
40
41
					       V a b -> a : f t
			     in f) t
42
43
44
\end{verbatim}
(note the NEED for shadowing)

45
46
=> Simplify

47
48
49
\begin{verbatim}
f = \ t -> case v of
	       V a b -> a : (letrec
50
				f = \ t -> a : f t
51
52
			     in f t)
\begin{verbatim}
53

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

57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
Other examples we'd like to catch with this kind of transformation

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

We'd like to avoid the redundant pattern match, transforming to

	last [] = error
	last (x:[]) = x
	last (x:(y:ys)) = last' y ys
		where
		  last' y []     = y
		  last' _ (y:ys) = last' y ys

	(is this necessarily an improvement)


Similarly drop:

	drop n [] = []
	drop 0 xs = xs
	drop n (x:xs) = drop (n-1) xs

Would like to pass n along unboxed.
	
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

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}
109
Here, the level of @f@ is zero, the level of @g@ is one,
110
111
112
and the level of @h@ is zero (NB not one).

\begin{code}
113
type LibCaseLevel = Int
114
115
116
117
118
119

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

\begin{code}
120
data LibCaseEnv
121
122
123
124
125
126
127
128
129
130
131
  = LibCaseEnv
	Int			-- Bomb-out size for deciding if
				-- potential liberatees are too big.
				-- (passed in from cmd-line args)

	LibCaseLevel		-- Current level

	(IdEnv LibCaseLevel)	-- Binds all non-top-level in-scope Ids
				-- (top-level and imported things have
				-- a level of zero)

132
	(IdEnv CoreBind)	-- Binds *only* recursively defined
133
134
135
136
137
138
139
140
141
142
143
144
145
				-- Ids, to their own binding group,
				-- and *only* in their own RHSs

	[(Id,LibCaseLevel)]     -- Each of these Ids was scrutinised by an
				-- enclosing case expression, with the
				-- specified number of enclosing
				-- recursive bindings; furthermore,
				-- the Id is bound at a lower level
				-- than the case expression.  The
				-- order is insignificant; it's a bag
				-- really

initEnv :: Int -> LibCaseEnv
146
initEnv bomb_size = LibCaseEnv bomb_size 0 emptyVarEnv emptyVarEnv []
147
148
149
150
151
152
153
154

bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
\end{code}


Programs
~~~~~~~~
\begin{code}
155
156
liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
liberateCase dflags binds
157
  = do {
158
	showPass dflags "Liberate case" ;
159
	let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
160
	endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
161
				{- no specific flag for dumping -} 
162
    }
163
164
165
166
167
168
169
170
171
172
173
  where
    do_prog env [] = []
    do_prog env (bind:binds) = bind' : do_prog env' binds
			     where
			       (env', bind') = libCaseBind env bind
\end{code}

Bindings
~~~~~~~~

\begin{code}
174
libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
175

176
177
libCaseBind env (NonRec binder rhs)
  = (addBinders env [binder], NonRec binder (libCase env rhs))
178

179
180
libCaseBind env (Rec pairs)
  = (env_body, Rec pairs')
181
182
183
184
185
186
187
188
189
190
191
192
  where
    (binders, rhss) = unzip pairs

    env_body = addBinders env binders

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

    env_rhs = if all rhs_small_enough rhss then extended_env else env

	-- 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!
193
194
	--
    extended_env = addRecBinds env [ (setIdNotExported binder, libCase env_body rhs)
195
				   | (binder, rhs) <- pairs ]
196

197
198
199
200
201
202
203
204
205
206
207
208
	-- Two subtle things: 
	-- (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.
    adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))

209
    rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
210
    lIBERATE_BOMB_SIZE   = bombOutSize env
211
212
213
214
215
216
217
218
\end{code}


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

\begin{code}
libCase :: LibCaseEnv
219
220
	-> CoreExpr
	-> CoreExpr
221

222
libCase env (Var v)		= libCaseId env v
223
libCase env (Lit lit)		= Lit lit
224
225
libCase env (Type ty)		= Type ty
libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
226
libCase env (Note note body)    = Note note (libCase env body)
227

228
229
libCase env (Lam binder body)
  = Lam binder (libCase (addBinders env [binder]) body)
230

231
232
libCase env (Let bind body)
  = Let bind' (libCase env_body body)
233
234
235
  where
    (env_body, bind') = libCaseBind env bind

236
237
libCase env (Case scrut bndr ty alts)
  = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
238
  where
239
    env_alts = addBinders env_with_scrut [bndr]
240
241
    env_with_scrut = case scrut of
		  	Var scrut_var -> addScrutedVar env scrut_var
242
			other	      -> env
243

244
libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
245
246
\end{code}

247
248
Ids
~~~
249
\begin{code}
250
libCaseId :: LibCaseEnv -> Id -> CoreExpr
251
libCaseId env v
252
  | Just the_bind <- lookupRecId env v	-- It's a use of a recursive thing
sof's avatar
sof committed
253
  , notNull free_scruts 		-- with free vars scrutinised in RHS
254
  = Let the_bind (Var v)
255
256

  | otherwise
257
  = Var v
258
259

  where
260
261
    rec_id_level = lookupLevel env v
    free_scruts  = freeScruts env rec_id_level
262
\end{code}
263
264
265
266
267
268



Utility functions
~~~~~~~~~~~~~~~~~
\begin{code}
269
addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
270
271
272
addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
  = LibCaseEnv bomb lvl lvl_env' rec_env scruts
  where
273
    lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
274

275
addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
276
277
278
279
addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
  = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
  where
    lvl'     = lvl + 1
280
281
    lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
    rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
282

283
addScrutedVar :: LibCaseEnv
284
	      -> Id		-- This Id is being scrutinised by a case expression
285
	      -> LibCaseEnv
286
287
288
289
290

addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
  | bind_lvl < lvl
  = LibCaseEnv bomb lvl lvl_env rec_env scruts'
	-- Add to scruts iff the scrut_var is being scrutinised at
291
	-- a deeper level than its defn
292
293
294
295

  | otherwise = env
  where
    scruts'  = (scrut_var, lvl) : scruts
296
    bind_lvl = case lookupVarEnv lvl_env scrut_var of
297
		 Just lvl -> lvl
298
		 Nothing  -> topLevel
299

300
lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
301
lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
302
  = lookupVarEnv rec_env id
303
304
305

lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
306
  = case lookupVarEnv lvl_env id of
307
      Just lvl -> lvl
308
      Nothing  -> topLevel
309

310
freeScruts :: LibCaseEnv
311
	   -> LibCaseLevel 	-- Level of the recursive Id
312
313
	   -> [Id]		-- Ids that are scrutinised between the binding
				-- of the recursive Id and here
314
freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
315
  = [v | (v,scrut_lvl) <- scruts, scrut_lvl > rec_bind_lvl]
316
\end{code}