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