Use Id.isStrictId
[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 module LiberateCase ( liberateCase ) where
8
9 #include "HsVersions.h"
10
11 import DynFlags
12 import HscTypes
13 import CoreLint         ( showPass, endPass )
14 import CoreSyn
15 import CoreUnfold       ( couldBeSmallEnoughToInline )
16 import Rules            ( RuleBase )
17 import UniqSupply       ( UniqSupply )
18 import SimplMonad       ( SimplCount, zeroSimplCount )
19 import Id
20 import FamInstEnv
21 import Type
22 import Coercion
23 import TyCon
24 import VarEnv
25 import Name             ( localiseName )
26 import Outputable
27 import Util             ( notNull )
28 import Data.IORef       ( readIORef )
29 \end{code}
30
31 The liberate-case transformation
32 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
33 This module walks over @Core@, and looks for @case@ on free variables.
34 The criterion is:
35         if there is case on a free on the route to the recursive call,
36         then the recursive call is replaced with an unfolding.
37
38 Example
39
40    f = \ t -> case v of
41                  V a b -> a : f t
42
43 => the inner f is replaced.
44
45    f = \ t -> case v of
46                  V a b -> a : (letrec
47                                 f =  \ t -> case v of
48                                                V a b -> a : f t
49                                in f) t
50 (note the NEED for shadowing)
51
52 => Simplify
53
54   f = \ t -> case v of
55                  V a b -> a : (letrec
56                                 f = \ t -> a : f t
57                                in f t)
58
59 Better code, because 'a' is  free inside the inner letrec, rather
60 than needing projection from v.
61
62 Other examples we'd like to catch with this kind of transformation
63
64         last []     = error 
65         last (x:[]) = x
66         last (x:xs) = last xs
67
68 We'd like to avoid the redundant pattern match, transforming to
69
70         last [] = error
71         last (x:[]) = x
72         last (x:(y:ys)) = last' y ys
73                 where
74                   last' y []     = y
75                   last' _ (y:ys) = last' y ys
76
77         (is this necessarily an improvement)
78
79 Similarly drop:
80
81         drop n [] = []
82         drop 0 xs = xs
83         drop n (x:xs) = drop (n-1) xs
84
85 Would like to pass n along unboxed.
86         
87 Note [Scrutinee with cast]
88 ~~~~~~~~~~~~~~~~~~~~~~~~~~
89 Consider this:
90     f = \ t -> case (v `cast` co) of
91                  V a b -> a : f t
92
93 Exactly the same optimisation (unrolling one call to f) will work here, 
94 despite the cast.  See mk_alt_env in the Case branch of libCase.
95
96
97 To think about (Apr 94)
98 ~~~~~~~~~~~~~~
99
100 Main worry: duplicating code excessively.  At the moment we duplicate
101 the entire binding group once at each recursive call.  But there may
102 be a group of recursive calls which share a common set of evaluated
103 free variables, in which case the duplication is a plain waste.
104
105 Another thing we could consider adding is some unfold-threshold thing,
106 so that we'll only duplicate if the size of the group rhss isn't too
107 big.
108
109 Data types
110 ~~~~~~~~~~
111 The ``level'' of a binder tells how many
112 recursive defns lexically enclose the binding
113 A recursive defn "encloses" its RHS, not its
114 scope.  For example:
115 \begin{verbatim}
116         letrec f = let g = ... in ...
117         in
118         let h = ...
119         in ...
120 \end{verbatim}
121 Here, the level of @f@ is zero, the level of @g@ is one,
122 and the level of @h@ is zero (NB not one).
123
124 Note [Indexed data types]
125 ~~~~~~~~~~~~~~~~~~~~~~~~~
126 Consider
127         data family T :: * -> *
128         data T Int = TI Int
129
130         f :: T Int -> Bool
131         f x = case x of { DEFAULT -> <body> }
132
133 We would like to change this to
134         f x = case x `cast` co of { TI p -> <body> }
135
136 so that <body> can make use of the fact that x is already evaluated to
137 a TI; and a case on a known data type may be more efficient than a
138 polymorphic one (not sure this is true any longer).  Anyway the former
139 showed up in Roman's experiments.  Example:
140   foo :: FooT Int -> Int -> Int
141   foo t n = t `seq` bar n
142      where
143        bar 0 = 0
144        bar n = bar (n - case t of TI i -> i)
145 Here we'd like to avoid repeated evaluating t inside the loop, by 
146 taking advantage of the `seq`.
147
148 We implement this as part of the liberate-case transformation by 
149 spotting
150         case <scrut> of (x::T) tys { DEFAULT ->  <body> }
151 where x :: T tys, and T is a indexed family tycon.  Find the
152 representation type (T77 tys'), and coercion co, and transform to
153         case <scrut> `cast` co of (y::T77 tys')
154             DEFAULT -> let x = y `cast` sym co in <body>
155
156 The "find the representation type" part is done by looking up in the
157 family-instance environment.
158
159 NB: in fact we re-use x (changing its type) to avoid making a fresh y;
160 this entails shadowing, but that's ok.
161
162 %************************************************************************
163 %*                                                                      *
164          Top-level code
165 %*                                                                      *
166 %************************************************************************
167
168 \begin{code}
169 liberateCase :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
170              -> IO (SimplCount, ModGuts)
171 liberateCase hsc_env _ _ guts
172   = do  { let dflags = hsc_dflags hsc_env
173         ; eps <- readIORef (hsc_EPS hsc_env)
174         ; let fam_envs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
175
176         ; showPass dflags "Liberate case"
177         ; let { env = initEnv dflags fam_envs
178               ; binds' = do_prog env (mg_binds guts) }
179         ; endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
180                         {- no specific flag for dumping -} 
181         ; return (zeroSimplCount dflags, guts { mg_binds = binds' }) }
182   where
183     do_prog env [] = []
184     do_prog env (bind:binds) = bind' : do_prog env' binds
185                              where
186                                (env', bind') = libCaseBind env bind
187 \end{code}
188
189
190 %************************************************************************
191 %*                                                                      *
192          Main payload
193 %*                                                                      *
194 %************************************************************************
195
196 Bindings
197 ~~~~~~~~
198 \begin{code}
199 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
200
201 libCaseBind env (NonRec binder rhs)
202   = (addBinders env [binder], NonRec binder (libCase env rhs))
203
204 libCaseBind env (Rec pairs)
205   = (env_body, Rec pairs')
206   where
207     (binders, rhss) = unzip pairs
208
209     env_body = addBinders env binders
210
211     pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
212
213     env_rhs = if all rhs_small_enough rhss then extended_env else env
214
215         -- We extend the rec-env by binding each Id to its rhs, first
216         -- processing the rhs with an *un-extended* environment, so
217         -- that the same process doesn't occur for ever!
218         --
219     extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
220                                    | (binder, rhs) <- pairs ]
221
222         -- Two subtle things: 
223         -- (a)  Reset the export flags on the binders so
224         --      that we don't get name clashes on exported things if the 
225         --      local binding floats out to top level.  This is most unlikely
226         --      to happen, since the whole point concerns free variables. 
227         --      But resetting the export flag is right regardless.
228         -- 
229         -- (b)  Make the name an Internal one.  External Names should never be
230         --      nested; if it were floated to the top level, we'd get a name
231         --      clash at code generation time.
232     adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
233
234     rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
235     lIBERATE_BOMB_SIZE   = bombOutSize env
236 \end{code}
237
238
239 Expressions
240 ~~~~~~~~~~~
241
242 \begin{code}
243 libCase :: LibCaseEnv
244         -> CoreExpr
245         -> CoreExpr
246
247 libCase env (Var v)             = libCaseId env v
248 libCase env (Lit lit)           = Lit lit
249 libCase env (Type ty)           = Type ty
250 libCase env (App fun arg)       = App (libCase env fun) (libCase env arg)
251 libCase env (Note note body)    = Note note (libCase env body)
252 libCase env (Cast e co)         = Cast (libCase env e) co
253
254 libCase env (Lam binder body)
255   = Lam binder (libCase (addBinders env [binder]) body)
256
257 libCase env (Let bind body)
258   = Let bind' (libCase env_body body)
259   where
260     (env_body, bind') = libCaseBind env bind
261
262 libCase env (Case scrut bndr ty alts)
263   = mkCase env (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
264   where
265     env_alts = addBinders (mk_alt_env scrut) [bndr]
266     mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
267     mk_alt_env (Cast scrut _)  = mk_alt_env scrut       -- Note [Scrutinee with cast]
268     mk_alt_env otehr           = env
269
270 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
271 \end{code}
272
273 \begin{code}
274 mkCase :: LibCaseEnv -> CoreExpr -> Id -> Type -> [CoreAlt] -> CoreExpr
275 -- See Note [Indexed data types]
276 mkCase env scrut bndr ty [(DEFAULT,_,rhs)]
277   | Just (tycon, tys)   <- splitTyConApp_maybe (idType bndr)
278   , [(subst, fam_inst)] <- lookupFamInstEnv (lc_fams env) tycon tys
279   = let 
280         rep_tc     = famInstTyCon fam_inst
281         rep_tys    = map (substTyVar subst) (tyConTyVars rep_tc)
282         bndr'      = setIdType bndr (mkTyConApp rep_tc rep_tys)
283         Just co_tc = tyConFamilyCoercion_maybe rep_tc
284         co         = mkTyConApp co_tc rep_tys
285         bind       = NonRec bndr (Cast (Var bndr') (mkSymCoercion co))
286     in mkCase env (Cast scrut co) bndr' ty [(DEFAULT,[],Let bind rhs)]
287 mkCase env scrut bndr ty alts
288   = Case scrut bndr ty alts
289 \end{code}
290
291 Ids
292 ~~~
293 \begin{code}
294 libCaseId :: LibCaseEnv -> Id -> CoreExpr
295 libCaseId env v
296   | Just the_bind <- lookupRecId env v  -- It's a use of a recursive thing
297   , notNull free_scruts                 -- with free vars scrutinised in RHS
298   = Let the_bind (Var v)
299
300   | otherwise
301   = Var v
302
303   where
304     rec_id_level = lookupLevel env v
305     free_scruts  = freeScruts env rec_id_level
306 \end{code}
307
308
309 %************************************************************************
310 %*                                                                      *
311         Utility functions
312 %*                                                                      *
313 %************************************************************************
314
315 \begin{code}
316 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
317 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
318   = env { lc_lvl_env = lvl_env' }
319   where
320     lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
321
322 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
323 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env, 
324                              lc_rec_env = rec_env}) pairs
325   = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
326   where
327     lvl'     = lvl + 1
328     lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
329     rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
330
331 addScrutedVar :: LibCaseEnv
332               -> Id             -- This Id is being scrutinised by a case expression
333               -> LibCaseEnv
334
335 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env, 
336                                 lc_scruts = scruts }) scrut_var
337   | bind_lvl < lvl
338   = env { lc_scruts = scruts' }
339         -- Add to scruts iff the scrut_var is being scrutinised at
340         -- a deeper level than its defn
341
342   | otherwise = env
343   where
344     scruts'  = (scrut_var, lvl) : scruts
345     bind_lvl = case lookupVarEnv lvl_env scrut_var of
346                  Just lvl -> lvl
347                  Nothing  -> topLevel
348
349 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
350 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
351
352 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
353 lookupLevel env id
354   = case lookupVarEnv (lc_lvl_env env) id of
355       Just lvl -> lc_lvl env
356       Nothing  -> topLevel
357
358 freeScruts :: LibCaseEnv
359            -> LibCaseLevel      -- Level of the recursive Id
360            -> [Id]              -- Ids that are scrutinised between the binding
361                                 -- of the recursive Id and here
362 freeScruts env rec_bind_lvl
363   = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
364 \end{code}
365
366 %************************************************************************
367 %*                                                                      *
368          The environment
369 %*                                                                      *
370 %************************************************************************
371
372 \begin{code}
373 type LibCaseLevel = Int
374
375 topLevel :: LibCaseLevel
376 topLevel = 0
377 \end{code}
378
379 \begin{code}
380 data LibCaseEnv
381   = LibCaseEnv {
382         lc_size :: Int,         -- Bomb-out size for deciding if
383                                 -- potential liberatees are too big.
384                                 -- (passed in from cmd-line args)
385
386         lc_lvl :: LibCaseLevel, -- Current level
387
388         lc_lvl_env :: IdEnv LibCaseLevel,  
389                         -- Binds all non-top-level in-scope Ids
390                         -- (top-level and imported things have
391                         -- a level of zero)
392
393         lc_rec_env :: IdEnv CoreBind, 
394                         -- Binds *only* recursively defined ids, 
395                         -- to their own binding group,
396                         -- and *only* in their own RHSs
397
398         lc_scruts :: [(Id,LibCaseLevel)],
399                         -- Each of these Ids was scrutinised by an
400                         -- enclosing case expression, with the
401                         -- specified number of enclosing
402                         -- recursive bindings; furthermore,
403                         -- the Id is bound at a lower level
404                         -- than the case expression.  The order is
405                         -- insignificant; it's a bag really
406
407         lc_fams :: FamInstEnvs
408                         -- Instance env for indexed data types 
409         }
410
411 initEnv :: DynFlags -> FamInstEnvs -> LibCaseEnv
412 initEnv dflags fams
413   = LibCaseEnv { lc_size = libCaseThreshold dflags,
414                  lc_lvl = 0,
415                  lc_lvl_env = emptyVarEnv, 
416                  lc_rec_env = emptyVarEnv,
417                  lc_scruts = [],
418                  lc_fams = fams }
419
420 bombOutSize = lc_size
421 \end{code}
422
423