2 % (c) The AQUA Project, Glasgow University, 1994-1998
4 \section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
7 module LiberateCase ( liberateCase ) where
9 #include "HsVersions.h"
13 import CoreUnfold ( couldBeSmallEnoughToInline )
16 import Util ( notNull )
19 The liberate-case transformation
20 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
21 This module walks over @Core@, and looks for @case@ on free variables.
23 if there is case on a free on the route to the recursive call,
24 then the recursive call is replaced with an unfolding.
31 => the inner f is replaced.
38 (note the NEED for shadowing)
47 Better code, because 'a' is free inside the inner letrec, rather
48 than needing projection from v.
50 Note that this deals with *free variables*. SpecConstr deals with
51 *arguments* that are of known form. E.g.
58 Note [Scrutinee with cast]
59 ~~~~~~~~~~~~~~~~~~~~~~~~~~
61 f = \ t -> case (v `cast` co) of
64 Exactly the same optimisation (unrolling one call to f) will work here,
65 despite the cast. See mk_alt_env in the Case branch of libCase.
68 Note [Only functions!]
69 ~~~~~~~~~~~~~~~~~~~~~~
70 Consider the following code
72 f = g (case v of V a b -> a : t f)
74 where g is expensive. If we aren't careful, liberate case will turn this into
77 V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
81 Yikes! We evaluate g twice. This leads to a O(2^n) explosion
82 if g calls back to the same code recursively.
84 Solution: make sure that we only do the liberate-case thing on *functions*
86 To think about (Apr 94)
88 Main worry: duplicating code excessively. At the moment we duplicate
89 the entire binding group once at each recursive call. But there may
90 be a group of recursive calls which share a common set of evaluated
91 free variables, in which case the duplication is a plain waste.
93 Another thing we could consider adding is some unfold-threshold thing,
94 so that we'll only duplicate if the size of the group rhss isn't too
99 The ``level'' of a binder tells how many
100 recursive defns lexically enclose the binding
101 A recursive defn "encloses" its RHS, not its
104 letrec f = let g = ... in ...
109 Here, the level of @f@ is zero, the level of @g@ is one,
110 and the level of @h@ is zero (NB not one).
113 %************************************************************************
117 %************************************************************************
120 liberateCase :: DynFlags -> [CoreBind] -> [CoreBind]
121 liberateCase dflags binds = do_prog (initEnv dflags) binds
124 do_prog env (bind:binds) = bind' : do_prog env' binds
126 (env', bind') = libCaseBind env bind
130 %************************************************************************
134 %************************************************************************
139 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
141 libCaseBind env (NonRec binder rhs)
142 = (addBinders env [binder], NonRec binder (libCase env rhs))
144 libCaseBind env (Rec pairs)
145 = (env_body, Rec pairs')
147 binders = map fst pairs
149 env_body = addBinders env binders
151 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
153 -- We extend the rec-env by binding each Id to its rhs, first
154 -- processing the rhs with an *un-extended* environment, so
155 -- that the same process doesn't occur for ever!
156 env_rhs = addRecBinds env [ (localiseId binder, libCase env_body rhs)
157 | (binder, rhs) <- pairs
158 , rhs_small_enough binder rhs ]
159 -- localiseID : see Note [Need to localiseId in libCaseBind]
162 rhs_small_enough id rhs -- Note [Small enough]
163 = idArity id > 0 -- Note [Only functions!]
164 && maybe True (\size -> couldBeSmallEnoughToInline size rhs)
168 Note [Need to localiseId in libCaseBind]
169 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
170 The call to localiseId is needed for two subtle reasons
171 (a) Reset the export flags on the binders so
172 that we don't get name clashes on exported things if the
173 local binding floats out to top level. This is most unlikely
174 to happen, since the whole point concerns free variables.
175 But resetting the export flag is right regardless.
177 (b) Make the name an Internal one. External Names should never be
178 nested; if it were floated to the top level, we'd get a name
179 clash at code generation time.
185 f = \x. BIG...(case fv of { (a,b) -> ...g.. })...
187 Then we *can* do liberate-case on g (small RHS) but not for f (too big).
188 But we can choose on a item-by-item basis, and that's what the
189 rhs_small_enough call in the comprehension for env_rhs does.
195 libCase :: LibCaseEnv
199 libCase env (Var v) = libCaseId env v
200 libCase _ (Lit lit) = Lit lit
201 libCase _ (Type ty) = Type ty
202 libCase _ (Coercion co) = Coercion co
203 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
204 libCase env (Note note body) = Note note (libCase env body)
205 libCase env (Cast e co) = Cast (libCase env e) co
207 libCase env (Lam binder body)
208 = Lam binder (libCase (addBinders env [binder]) body)
210 libCase env (Let bind body)
211 = Let bind' (libCase env_body body)
213 (env_body, bind') = libCaseBind env bind
215 libCase env (Case scrut bndr ty alts)
216 = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
218 env_alts = addBinders (mk_alt_env scrut) [bndr]
219 mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
220 mk_alt_env (Cast scrut _) = mk_alt_env scrut -- Note [Scrutinee with cast]
223 libCaseAlt :: LibCaseEnv -> (AltCon, [CoreBndr], CoreExpr)
224 -> (AltCon, [CoreBndr], CoreExpr)
225 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
232 libCaseId :: LibCaseEnv -> Id -> CoreExpr
234 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
235 , notNull free_scruts -- with free vars scrutinised in RHS
236 = Let the_bind (Var v)
242 rec_id_level = lookupLevel env v
243 free_scruts = freeScruts env rec_id_level
245 freeScruts :: LibCaseEnv
246 -> LibCaseLevel -- Level of the recursive Id
247 -> [Id] -- Ids that are scrutinised between the binding
248 -- of the recursive Id and here
249 freeScruts env rec_bind_lvl
250 = [v | (v, scrut_bind_lvl, scrut_at_lvl) <- lc_scruts env
251 , scrut_bind_lvl <= rec_bind_lvl
252 , scrut_at_lvl > rec_bind_lvl]
253 -- Note [When to specialise]
254 -- Note [Avoiding fruitless liberate-case]
257 Note [When to specialise]
258 ~~~~~~~~~~~~~~~~~~~~~~~~~
260 f = \x. letrec g = \y. case x of
261 True -> ... (f a) ...
262 False -> ... (g b) ...
264 We get the following levels
270 Then 'x' is being scrutinised at a deeper level than its binding, so
271 it's added to lc_sruts: [(x,1)]
273 We do *not* want to specialise the call to 'f', becuase 'x' is not free
274 in 'f'. So here the bind-level of 'x' (=1) is not <= the bind-level of 'f' (=0).
276 We *do* want to specialise the call to 'g', because 'x' is free in g.
277 Here the bind-level of 'x' (=1) is <= the bind-level of 'g' (=1).
279 Note [Avoiding fruitless liberate-case]
280 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
282 f = \x. case top_lvl_thing of
283 I# _ -> let g = \y. ... g ...
286 Here, top_lvl_thing is scrutinised at a level (1) deeper than its
287 binding site (0). Nevertheless, we do NOT want to specialise the call
288 to 'g' because all the structure in its free variables is already
289 visible at the definition site for g. Hence, when considering specialising
290 an occurrence of 'g', we want to check that there's a scruted-var v st
292 a) v's binding site is *outside* g
293 b) v's scrutinisation site is *inside* g
296 %************************************************************************
300 %************************************************************************
303 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
304 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
305 = env { lc_lvl_env = lvl_env' }
307 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
309 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
310 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
311 lc_rec_env = rec_env}) pairs
312 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
315 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
316 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
318 addScrutedVar :: LibCaseEnv
319 -> Id -- This Id is being scrutinised by a case expression
322 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
323 lc_scruts = scruts }) scrut_var
325 = env { lc_scruts = scruts' }
326 -- Add to scruts iff the scrut_var is being scrutinised at
327 -- a deeper level than its defn
331 scruts' = (scrut_var, bind_lvl, lvl) : scruts
332 bind_lvl = case lookupVarEnv lvl_env scrut_var of
336 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
337 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
339 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
341 = case lookupVarEnv (lc_lvl_env env) id of
346 %************************************************************************
350 %************************************************************************
353 type LibCaseLevel = Int
355 topLevel :: LibCaseLevel
362 lc_size :: Maybe Int, -- Bomb-out size for deciding if
363 -- potential liberatees are too big.
364 -- (passed in from cmd-line args)
366 lc_lvl :: LibCaseLevel, -- Current level
367 -- The level is incremented when (and only when) going
368 -- inside the RHS of a (sufficiently small) recursive
371 lc_lvl_env :: IdEnv LibCaseLevel,
372 -- Binds all non-top-level in-scope Ids (top-level and
373 -- imported things have a level of zero)
375 lc_rec_env :: IdEnv CoreBind,
376 -- Binds *only* recursively defined ids, to their own
377 -- binding group, and *only* in their own RHSs
379 lc_scruts :: [(Id, LibCaseLevel, LibCaseLevel)]
380 -- Each of these Ids was scrutinised by an enclosing
381 -- case expression, at a level deeper than its binding
384 -- The first LibCaseLevel is the *binding level* of
385 -- the scrutinised Id,
386 -- The second is the level *at which it was scrutinised*.
387 -- (see Note [Avoiding fruitless liberate-case])
388 -- The former is a bit redundant, since you could always
389 -- look it up in lc_lvl_env, but it's just cached here
391 -- The order is insignificant; it's a bag really
393 -- There's one element per scrutinisation;
394 -- in principle the same Id may appear multiple times,
395 -- although that'd be unusual:
396 -- case x of { (a,b) -> ....(case x of ...) .. }
399 initEnv :: DynFlags -> LibCaseEnv
401 = LibCaseEnv { lc_size = liberateCaseThreshold dflags,
403 lc_lvl_env = emptyVarEnv,
404 lc_rec_env = emptyVarEnv,
407 bombOutSize :: LibCaseEnv -> Maybe Int
408 bombOutSize = lc_size