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 CoreLint ( showPass, endPass )
15 import CoreUnfold ( couldBeSmallEnoughToInline )
16 import Rules ( RuleBase )
17 import UniqSupply ( UniqSupply )
18 import SimplMonad ( SimplCount, zeroSimplCount )
21 import Name ( localiseName )
22 import Util ( notNull )
25 The liberate-case transformation
26 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
27 This module walks over @Core@, and looks for @case@ on free variables.
29 if there is case on a free on the route to the recursive call,
30 then the recursive call is replaced with an unfolding.
37 => the inner f is replaced.
44 (note the NEED for shadowing)
53 Better code, because 'a' is free inside the inner letrec, rather
54 than needing projection from v.
56 Other examples we'd like to catch with this kind of transformation
62 We'd like to avoid the redundant pattern match, transforming to
66 last (x:(y:ys)) = last' y ys
69 last' _ (y:ys) = last' y ys
71 (is this necessarily an improvement)
77 drop n (x:xs) = drop (n-1) xs
79 Would like to pass n along unboxed.
81 Note [Scrutinee with cast]
82 ~~~~~~~~~~~~~~~~~~~~~~~~~~
84 f = \ t -> case (v `cast` co) of
87 Exactly the same optimisation (unrolling one call to f) will work here,
88 despite the cast. See mk_alt_env in the Case branch of libCase.
91 Note [Only functions!]
92 ~~~~~~~~~~~~~~~~~~~~~~
93 Consider the following code
95 f = g (case v of V a b -> a : t f)
97 where g is expensive. If we aren't careful, liberate case will turn this into
100 V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
104 Yikes! We evaluate g twice. This leads to a O(2^n) explosion
105 if g calls back to the same code recursively.
107 Solution: make sure that we only do the liberate-case thing on *functions*
109 To think about (Apr 94)
111 Main worry: duplicating code excessively. At the moment we duplicate
112 the entire binding group once at each recursive call. But there may
113 be a group of recursive calls which share a common set of evaluated
114 free variables, in which case the duplication is a plain waste.
116 Another thing we could consider adding is some unfold-threshold thing,
117 so that we'll only duplicate if the size of the group rhss isn't too
122 The ``level'' of a binder tells how many
123 recursive defns lexically enclose the binding
124 A recursive defn "encloses" its RHS, not its
127 letrec f = let g = ... in ...
132 Here, the level of @f@ is zero, the level of @g@ is one,
133 and the level of @h@ is zero (NB not one).
136 %************************************************************************
140 %************************************************************************
143 liberateCase :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
144 -> IO (SimplCount, ModGuts)
145 liberateCase hsc_env _ _ guts
146 = do { let dflags = hsc_dflags hsc_env
148 ; showPass dflags "Liberate case"
149 ; let { env = initEnv dflags
150 ; binds' = do_prog env (mg_binds guts) }
151 ; endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
152 {- no specific flag for dumping -}
153 ; return (zeroSimplCount dflags, guts { mg_binds = binds' }) }
156 do_prog env (bind:binds) = bind' : do_prog env' binds
158 (env', bind') = libCaseBind env bind
162 %************************************************************************
166 %************************************************************************
171 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
173 libCaseBind env (NonRec binder rhs)
174 = (addBinders env [binder], NonRec binder (libCase env rhs))
176 libCaseBind env (Rec pairs)
177 = (env_body, Rec pairs')
179 (binders, rhss) = unzip pairs
181 env_body = addBinders env binders
183 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
185 env_rhs = if all rhs_small_enough pairs then extended_env else env
187 -- We extend the rec-env by binding each Id to its rhs, first
188 -- processing the rhs with an *un-extended* environment, so
189 -- that the same process doesn't occur for ever!
191 extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
192 | (binder, rhs) <- pairs ]
194 -- Two subtle things:
195 -- (a) Reset the export flags on the binders so
196 -- that we don't get name clashes on exported things if the
197 -- local binding floats out to top level. This is most unlikely
198 -- to happen, since the whole point concerns free variables.
199 -- But resetting the export flag is right regardless.
201 -- (b) Make the name an Internal one. External Names should never be
202 -- nested; if it were floated to the top level, we'd get a name
203 -- clash at code generation time.
204 adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
206 rhs_small_enough (id,rhs)
207 = idArity id > 0 -- Note [Only functions!]
208 && couldBeSmallEnoughToInline (bombOutSize env) rhs
216 libCase :: LibCaseEnv
220 libCase env (Var v) = libCaseId env v
221 libCase env (Lit lit) = Lit lit
222 libCase env (Type ty) = Type ty
223 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
224 libCase env (Note note body) = Note note (libCase env body)
225 libCase env (Cast e co) = Cast (libCase env e) co
227 libCase env (Lam binder body)
228 = Lam binder (libCase (addBinders env [binder]) body)
230 libCase env (Let bind body)
231 = Let bind' (libCase env_body body)
233 (env_body, bind') = libCaseBind env bind
235 libCase env (Case scrut bndr ty alts)
236 = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
238 env_alts = addBinders (mk_alt_env scrut) [bndr]
239 mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
240 mk_alt_env (Cast scrut _) = mk_alt_env scrut -- Note [Scrutinee with cast]
241 mk_alt_env otehr = env
243 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
250 libCaseId :: LibCaseEnv -> Id -> CoreExpr
252 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
253 , notNull free_scruts -- with free vars scrutinised in RHS
254 = Let the_bind (Var v)
260 rec_id_level = lookupLevel env v
261 free_scruts = freeScruts env rec_id_level
265 %************************************************************************
269 %************************************************************************
272 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
273 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
274 = env { lc_lvl_env = lvl_env' }
276 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
278 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
279 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
280 lc_rec_env = rec_env}) pairs
281 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
284 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
285 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
287 addScrutedVar :: LibCaseEnv
288 -> Id -- This Id is being scrutinised by a case expression
291 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
292 lc_scruts = scruts }) scrut_var
294 = env { lc_scruts = scruts' }
295 -- Add to scruts iff the scrut_var is being scrutinised at
296 -- a deeper level than its defn
300 scruts' = (scrut_var, lvl) : scruts
301 bind_lvl = case lookupVarEnv lvl_env scrut_var of
305 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
306 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
308 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
310 = case lookupVarEnv (lc_lvl_env env) id of
314 freeScruts :: LibCaseEnv
315 -> LibCaseLevel -- Level of the recursive Id
316 -> [Id] -- Ids that are scrutinised between the binding
317 -- of the recursive Id and here
318 freeScruts env rec_bind_lvl
319 = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
322 %************************************************************************
326 %************************************************************************
329 type LibCaseLevel = Int
331 topLevel :: LibCaseLevel
338 lc_size :: Int, -- Bomb-out size for deciding if
339 -- potential liberatees are too big.
340 -- (passed in from cmd-line args)
342 lc_lvl :: LibCaseLevel, -- Current level
344 lc_lvl_env :: IdEnv LibCaseLevel,
345 -- Binds all non-top-level in-scope Ids
346 -- (top-level and imported things have
349 lc_rec_env :: IdEnv CoreBind,
350 -- Binds *only* recursively defined ids,
351 -- to their own binding group,
352 -- and *only* in their own RHSs
354 lc_scruts :: [(Id,LibCaseLevel)]
355 -- Each of these Ids was scrutinised by an
356 -- enclosing case expression, with the
357 -- specified number of enclosing
358 -- recursive bindings; furthermore,
359 -- the Id is bound at a lower level
360 -- than the case expression. The order is
361 -- insignificant; it's a bag really
364 initEnv :: DynFlags -> LibCaseEnv
366 = LibCaseEnv { lc_size = specThreshold dflags,
368 lc_lvl_env = emptyVarEnv,
369 lc_rec_env = emptyVarEnv,
372 bombOutSize = lc_size