never executed always true always false
1 {-
2 (c) The AQUA Project, Glasgow University, 1994-1998
3
4 \section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
5 -}
6
7
8 module GHC.Core.Opt.LiberateCase ( liberateCase ) where
9
10 import GHC.Prelude
11
12 import GHC.Driver.Session
13 import GHC.Core
14 import GHC.Core.Unfold
15 import GHC.Builtin.Types ( unitDataConId )
16 import GHC.Types.Id
17 import GHC.Types.Var.Env
18 import GHC.Utils.Misc ( notNull )
19
20 {-
21 The liberate-case transformation
22 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
23 This module walks over @Core@, and looks for @case@ on free variables.
24 The criterion is:
25 if there is case on a free on the route to the recursive call,
26 then the recursive call is replaced with an unfolding.
27
28 Example
29
30 f = \ t -> case v of
31 V a b -> a : f t
32
33 => the inner f is replaced.
34
35 f = \ t -> case v of
36 V a b -> a : (letrec
37 f = \ t -> case v of
38 V a b -> a : f t
39 in f) t
40 (note the NEED for shadowing)
41
42 => Simplify
43
44 f = \ t -> case v of
45 V a b -> a : (letrec
46 f = \ t -> a : f t
47 in f t)
48
49 Better code, because 'a' is free inside the inner letrec, rather
50 than needing projection from v.
51
52 Note that this deals with *free variables*. SpecConstr deals with
53 *arguments* that are of known form. E.g.
54
55 last [] = error
56 last (x:[]) = x
57 last (x:xs) = last xs
58
59
60 Note [Scrutinee with cast]
61 ~~~~~~~~~~~~~~~~~~~~~~~~~~
62 Consider this:
63 f = \ t -> case (v `cast` co) of
64 V a b -> a : f t
65
66 Exactly the same optimisation (unrolling one call to f) will work here,
67 despite the cast. See mk_alt_env in the Case branch of libCase.
68
69
70 To think about (Apr 94)
71 ~~~~~~~~~~~~~~
72 Main worry: duplicating code excessively. At the moment we duplicate
73 the entire binding group once at each recursive call. But there may
74 be a group of recursive calls which share a common set of evaluated
75 free variables, in which case the duplication is a plain waste.
76
77 Another thing we could consider adding is some unfold-threshold thing,
78 so that we'll only duplicate if the size of the group rhss isn't too
79 big.
80
81 Data types
82 ~~~~~~~~~~
83 The ``level'' of a binder tells how many
84 recursive defns lexically enclose the binding
85 A recursive defn "encloses" its RHS, not its
86 scope. For example:
87 \begin{verbatim}
88 letrec f = let g = ... in ...
89 in
90 let h = ...
91 in ...
92 \end{verbatim}
93 Here, the level of @f@ is zero, the level of @g@ is one,
94 and the level of @h@ is zero (NB not one).
95
96
97 ************************************************************************
98 * *
99 Top-level code
100 * *
101 ************************************************************************
102 -}
103
104 liberateCase :: DynFlags -> CoreProgram -> CoreProgram
105 liberateCase dflags binds = do_prog (initLiberateCaseEnv dflags) binds
106 where
107 do_prog _ [] = []
108 do_prog env (bind:binds) = bind' : do_prog env' binds
109 where
110 (env', bind') = libCaseBind env bind
111
112
113 initLiberateCaseEnv :: DynFlags -> LibCaseEnv
114 initLiberateCaseEnv dflags = LibCaseEnv
115 { lc_threshold = liberateCaseThreshold dflags
116 , lc_uf_opts = unfoldingOpts dflags
117 , lc_lvl = 0
118 , lc_lvl_env = emptyVarEnv
119 , lc_rec_env = emptyVarEnv
120 , lc_scruts = []
121 }
122
123 {-
124 ************************************************************************
125 * *
126 Main payload
127 * *
128 ************************************************************************
129
130 Bindings
131 ~~~~~~~~
132 -}
133
134 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
135
136 libCaseBind env (NonRec binder rhs)
137 = (addBinders env [binder], NonRec binder (libCase env rhs))
138
139 libCaseBind env (Rec pairs)
140 = (env_body, Rec pairs')
141 where
142 binders = map fst pairs
143
144 env_body = addBinders env binders
145
146 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
147
148 -- We extend the rec-env by binding each Id to its rhs, first
149 -- processing the rhs with an *un-extended* environment, so
150 -- that the same process doesn't occur for ever!
151 env_rhs | is_dupable_bind = addRecBinds env dup_pairs
152 | otherwise = env
153
154 dup_pairs = [ (localiseId binder, libCase env_body rhs)
155 | (binder, rhs) <- pairs ]
156 -- localiseID : see Note [Need to localiseId in libCaseBind]
157
158 is_dupable_bind = small_enough && all ok_pair pairs
159
160 -- Size: we are going to duplicate dup_pairs; to find their
161 -- size, build a fake binding (let { dup_pairs } in (),
162 -- and find the size of that
163 -- See Note [Small enough]
164 small_enough = case lc_threshold env of
165 Nothing -> True -- Infinity
166 Just size -> couldBeSmallEnoughToInline (lc_uf_opts env) size $
167 Let (Rec dup_pairs) (Var unitDataConId)
168
169 ok_pair (id,_)
170 = idArity id > 0 -- Note [Only functions!]
171 && not (isDeadEndId id) -- Note [Not bottoming ids]
172
173 {- Note [Not bottoming Ids]
174 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
175 Do not specialise error-functions (this is unusual, but I once saw it,
176 (actually in Data.Typable.Internal)
177
178 Note [Only functions!]
179 ~~~~~~~~~~~~~~~~~~~~~~
180 Consider the following code
181
182 f = g (case v of V a b -> a : t f)
183
184 where g is expensive. If we aren't careful, liberate case will turn this into
185
186 f = g (case v of
187 V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
188 in f)
189 )
190
191 Yikes! We evaluate g twice. This leads to a O(2^n) explosion
192 if g calls back to the same code recursively.
193
194 Solution: make sure that we only do the liberate-case thing on *functions*
195
196 Note [Small enough]
197 ~~~~~~~~~~~~~~~~~~~
198 Consider
199 \fv. letrec
200 f = \x. BIG...(case fv of { (a,b) -> ...g.. })...
201 g = \y. SMALL...f...
202
203 Then we *can* in principle do liberate-case on 'g' (small RHS) but not
204 for 'f' (too big). But doing so is not profitable, because duplicating
205 'g' at its call site in 'f' doesn't get rid of any cases. So we just
206 ask for the whole group to be small enough.
207
208 Note [Need to localiseId in libCaseBind]
209 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
210 The call to localiseId is needed for two subtle reasons
211 (a) Reset the export flags on the binders so
212 that we don't get name clashes on exported things if the
213 local binding floats out to top level. This is most unlikely
214 to happen, since the whole point concerns free variables.
215 But resetting the export flag is right regardless.
216
217 (b) Make the name an Internal one. External Names should never be
218 nested; if it were floated to the top level, we'd get a name
219 clash at code generation time.
220
221 Expressions
222 ~~~~~~~~~~~
223 -}
224
225 libCase :: LibCaseEnv
226 -> CoreExpr
227 -> CoreExpr
228
229 libCase env (Var v) = libCaseApp env v []
230 libCase _ (Lit lit) = Lit lit
231 libCase _ (Type ty) = Type ty
232 libCase _ (Coercion co) = Coercion co
233 libCase env e@(App {}) | let (fun, args) = collectArgs e
234 , Var v <- fun
235 = libCaseApp env v args
236 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
237 libCase env (Tick tickish body) = Tick tickish (libCase env body)
238 libCase env (Cast e co) = Cast (libCase env e) co
239
240 libCase env (Lam binder body)
241 = Lam binder (libCase (addBinders env [binder]) body)
242
243 libCase env (Let bind body)
244 = Let bind' (libCase env_body body)
245 where
246 (env_body, bind') = libCaseBind env bind
247
248 libCase env (Case scrut bndr ty alts)
249 = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
250 where
251 env_alts = addBinders (mk_alt_env scrut) [bndr]
252 mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
253 mk_alt_env (Cast scrut _) = mk_alt_env scrut -- Note [Scrutinee with cast]
254 mk_alt_env _ = env
255
256 libCaseAlt :: LibCaseEnv -> Alt CoreBndr -> Alt CoreBndr
257 libCaseAlt env (Alt con args rhs) = Alt con args (libCase (addBinders env args) rhs)
258
259 {-
260 Ids
261 ~~~
262
263 To unfold, we can't just wrap the id itself in its binding if it's a join point:
264
265 jump j a b c => (joinrec j x y z = ... in jump j) a b c -- wrong!!!
266
267 Every jump must provide all arguments, so we have to be careful to wrap the
268 whole jump instead:
269
270 jump j a b c => joinrec j x y z = ... in jump j a b c -- right
271
272 -}
273
274 libCaseApp :: LibCaseEnv -> Id -> [CoreExpr] -> CoreExpr
275 libCaseApp env v args
276 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
277 , notNull free_scruts -- with free vars scrutinised in RHS
278 = Let the_bind expr'
279
280 | otherwise
281 = expr'
282
283 where
284 rec_id_level = lookupLevel env v
285 free_scruts = freeScruts env rec_id_level
286 expr' = mkApps (Var v) (map (libCase env) args)
287
288 freeScruts :: LibCaseEnv
289 -> LibCaseLevel -- Level of the recursive Id
290 -> [Id] -- Ids that are scrutinised between the binding
291 -- of the recursive Id and here
292 freeScruts env rec_bind_lvl
293 = [v | (v, scrut_bind_lvl, scrut_at_lvl) <- lc_scruts env
294 , scrut_bind_lvl <= rec_bind_lvl
295 , scrut_at_lvl > rec_bind_lvl]
296 -- Note [When to specialise]
297 -- Note [Avoiding fruitless liberate-case]
298
299 {-
300 Note [When to specialise]
301 ~~~~~~~~~~~~~~~~~~~~~~~~~
302 Consider
303 f = \x. letrec g = \y. case x of
304 True -> ... (f a) ...
305 False -> ... (g b) ...
306
307 We get the following levels
308 f 0
309 x 1
310 g 1
311 y 2
312
313 Then 'x' is being scrutinised at a deeper level than its binding, so
314 it's added to lc_sruts: [(x,1)]
315
316 We do *not* want to specialise the call to 'f', because 'x' is not free
317 in 'f'. So here the bind-level of 'x' (=1) is not <= the bind-level of 'f' (=0).
318
319 We *do* want to specialise the call to 'g', because 'x' is free in g.
320 Here the bind-level of 'x' (=1) is <= the bind-level of 'g' (=1).
321
322 Note [Avoiding fruitless liberate-case]
323 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
324 Consider also:
325 f = \x. case top_lvl_thing of
326 I# _ -> let g = \y. ... g ...
327 in ...
328
329 Here, top_lvl_thing is scrutinised at a level (1) deeper than its
330 binding site (0). Nevertheless, we do NOT want to specialise the call
331 to 'g' because all the structure in its free variables is already
332 visible at the definition site for g. Hence, when considering specialising
333 an occurrence of 'g', we want to check that there's a scruted-var v st
334
335 a) v's binding site is *outside* g
336 b) v's scrutinisation site is *inside* g
337
338
339 ************************************************************************
340 * *
341 Utility functions
342 * *
343 ************************************************************************
344 -}
345
346 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
347 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
348 = env { lc_lvl_env = lvl_env' }
349 where
350 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
351
352 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
353 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
354 lc_rec_env = rec_env}) pairs
355 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
356 where
357 lvl' = lvl + 1
358 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
359 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
360
361 addScrutedVar :: LibCaseEnv
362 -> Id -- This Id is being scrutinised by a case expression
363 -> LibCaseEnv
364
365 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
366 lc_scruts = scruts }) scrut_var
367 | bind_lvl < lvl
368 = env { lc_scruts = scruts' }
369 -- Add to scruts iff the scrut_var is being scrutinised at
370 -- a deeper level than its defn
371
372 | otherwise = env
373 where
374 scruts' = (scrut_var, bind_lvl, lvl) : scruts
375 bind_lvl = case lookupVarEnv lvl_env scrut_var of
376 Just lvl -> lvl
377 Nothing -> topLevel
378
379 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
380 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
381
382 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
383 lookupLevel env id
384 = case lookupVarEnv (lc_lvl_env env) id of
385 Just lvl -> lvl
386 Nothing -> topLevel
387
388 {-
389 ************************************************************************
390 * *
391 The environment
392 * *
393 ************************************************************************
394 -}
395
396 type LibCaseLevel = Int
397
398 topLevel :: LibCaseLevel
399 topLevel = 0
400
401 data LibCaseEnv
402 = LibCaseEnv {
403 lc_threshold :: Maybe Int,
404 -- ^ Bomb-out size for deciding if potential liberatees are too
405 -- big.
406
407 lc_uf_opts :: UnfoldingOpts,
408 -- ^ Unfolding options
409
410 lc_lvl :: LibCaseLevel, -- ^ Current level
411 -- The level is incremented when (and only when) going
412 -- inside the RHS of a (sufficiently small) recursive
413 -- function.
414
415 lc_lvl_env :: IdEnv LibCaseLevel,
416 -- ^ Binds all non-top-level in-scope Ids (top-level and
417 -- imported things have a level of zero)
418
419 lc_rec_env :: IdEnv CoreBind,
420 -- ^ Binds *only* recursively defined ids, to their own
421 -- binding group, and *only* in their own RHSs
422
423 lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
424 -- ^ Each of these Ids was scrutinised by an enclosing
425 -- case expression, at a level deeper than its binding
426 -- level.
427 --
428 -- The first LibCaseLevel is the *binding level* of
429 -- the scrutinised Id,
430 -- The second is the level *at which it was scrutinised*.
431 -- (see Note [Avoiding fruitless liberate-case])
432 -- The former is a bit redundant, since you could always
433 -- look it up in lc_lvl_env, but it's just cached here
434 --
435 -- The order is insignificant; it's a bag really
436 --
437 -- There's one element per scrutinisation;
438 -- in principle the same Id may appear multiple times,
439 -- although that'd be unusual:
440 -- case x of { (a,b) -> ....(case x of ...) .. }
441 }