2 % (c) The AQUA Project, Glasgow University, 1994-1998
4 \section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
8 -- The above warning supression flag is a temporary kludge.
9 -- While working on this module you are encouraged to remove it and fix
10 -- any warnings in the module. See
11 -- http://hackage.haskell.org/trac/ghc/wiki/WorkingConventions#Warnings
14 module LiberateCase ( liberateCase ) where
16 #include "HsVersions.h"
20 import CoreLint ( showPass, endPass )
22 import CoreUnfold ( couldBeSmallEnoughToInline )
23 import Rules ( RuleBase )
24 import UniqSupply ( UniqSupply )
25 import SimplMonad ( SimplCount, zeroSimplCount )
28 import Name ( localiseName )
29 import Util ( notNull )
32 The liberate-case transformation
33 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
34 This module walks over @Core@, and looks for @case@ on free variables.
36 if there is case on a free on the route to the recursive call,
37 then the recursive call is replaced with an unfolding.
44 => the inner f is replaced.
51 (note the NEED for shadowing)
60 Better code, because 'a' is free inside the inner letrec, rather
61 than needing projection from v.
63 Other examples we'd like to catch with this kind of transformation
69 We'd like to avoid the redundant pattern match, transforming to
73 last (x:(y:ys)) = last' y ys
76 last' _ (y:ys) = last' y ys
78 (is this necessarily an improvement)
84 drop n (x:xs) = drop (n-1) xs
86 Would like to pass n along unboxed.
88 Note [Scrutinee with cast]
89 ~~~~~~~~~~~~~~~~~~~~~~~~~~
91 f = \ t -> case (v `cast` co) of
94 Exactly the same optimisation (unrolling one call to f) will work here,
95 despite the cast. See mk_alt_env in the Case branch of libCase.
98 Note [Only functions!]
99 ~~~~~~~~~~~~~~~~~~~~~~
100 Consider the following code
102 f = g (case v of V a b -> a : t f)
104 where g is expensive. If we aren't careful, liberate case will turn this into
107 V a b -> a : t (letrec f = g (case v of V a b -> a : f t)
111 Yikes! We evaluate g twice. This leads to a O(2^n) explosion
112 if g calls back to the same code recursively.
114 Solution: make sure that we only do the liberate-case thing on *functions*
116 To think about (Apr 94)
118 Main worry: duplicating code excessively. At the moment we duplicate
119 the entire binding group once at each recursive call. But there may
120 be a group of recursive calls which share a common set of evaluated
121 free variables, in which case the duplication is a plain waste.
123 Another thing we could consider adding is some unfold-threshold thing,
124 so that we'll only duplicate if the size of the group rhss isn't too
129 The ``level'' of a binder tells how many
130 recursive defns lexically enclose the binding
131 A recursive defn "encloses" its RHS, not its
134 letrec f = let g = ... in ...
139 Here, the level of @f@ is zero, the level of @g@ is one,
140 and the level of @h@ is zero (NB not one).
143 %************************************************************************
147 %************************************************************************
150 liberateCase :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
151 -> IO (SimplCount, ModGuts)
152 liberateCase hsc_env _ _ guts
153 = do { let dflags = hsc_dflags hsc_env
155 ; showPass dflags "Liberate case"
156 ; let { env = initEnv dflags
157 ; binds' = do_prog env (mg_binds guts) }
158 ; endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
159 {- no specific flag for dumping -}
160 ; return (zeroSimplCount dflags, guts { mg_binds = binds' }) }
163 do_prog env (bind:binds) = bind' : do_prog env' binds
165 (env', bind') = libCaseBind env bind
169 %************************************************************************
173 %************************************************************************
178 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
180 libCaseBind env (NonRec binder rhs)
181 = (addBinders env [binder], NonRec binder (libCase env rhs))
183 libCaseBind env (Rec pairs)
184 = (env_body, Rec pairs')
186 (binders, rhss) = unzip pairs
188 env_body = addBinders env binders
190 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
192 env_rhs = if all rhs_small_enough pairs then extended_env else env
194 -- We extend the rec-env by binding each Id to its rhs, first
195 -- processing the rhs with an *un-extended* environment, so
196 -- that the same process doesn't occur for ever!
198 extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
199 | (binder, rhs) <- pairs ]
201 -- Two subtle things:
202 -- (a) Reset the export flags on the binders so
203 -- that we don't get name clashes on exported things if the
204 -- local binding floats out to top level. This is most unlikely
205 -- to happen, since the whole point concerns free variables.
206 -- But resetting the export flag is right regardless.
208 -- (b) Make the name an Internal one. External Names should never be
209 -- nested; if it were floated to the top level, we'd get a name
210 -- clash at code generation time.
211 adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
213 rhs_small_enough (id,rhs)
214 = idArity id > 0 -- Note [Only functions!]
215 && couldBeSmallEnoughToInline (bombOutSize env) rhs
223 libCase :: LibCaseEnv
227 libCase env (Var v) = libCaseId env v
228 libCase env (Lit lit) = Lit lit
229 libCase env (Type ty) = Type ty
230 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
231 libCase env (Note note body) = Note note (libCase env body)
232 libCase env (Cast e co) = Cast (libCase env e) co
234 libCase env (Lam binder body)
235 = Lam binder (libCase (addBinders env [binder]) body)
237 libCase env (Let bind body)
238 = Let bind' (libCase env_body body)
240 (env_body, bind') = libCaseBind env bind
242 libCase env (Case scrut bndr ty alts)
243 = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
245 env_alts = addBinders (mk_alt_env scrut) [bndr]
246 mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
247 mk_alt_env (Cast scrut _) = mk_alt_env scrut -- Note [Scrutinee with cast]
248 mk_alt_env otehr = env
250 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
257 libCaseId :: LibCaseEnv -> Id -> CoreExpr
259 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
260 , notNull free_scruts -- with free vars scrutinised in RHS
261 = Let the_bind (Var v)
267 rec_id_level = lookupLevel env v
268 free_scruts = freeScruts env rec_id_level
272 %************************************************************************
276 %************************************************************************
279 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
280 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
281 = env { lc_lvl_env = lvl_env' }
283 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
285 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
286 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
287 lc_rec_env = rec_env}) pairs
288 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
291 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
292 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
294 addScrutedVar :: LibCaseEnv
295 -> Id -- This Id is being scrutinised by a case expression
298 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
299 lc_scruts = scruts }) scrut_var
301 = env { lc_scruts = scruts' }
302 -- Add to scruts iff the scrut_var is being scrutinised at
303 -- a deeper level than its defn
307 scruts' = (scrut_var, lvl) : scruts
308 bind_lvl = case lookupVarEnv lvl_env scrut_var of
312 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
313 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
315 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
317 = case lookupVarEnv (lc_lvl_env env) id of
321 freeScruts :: LibCaseEnv
322 -> LibCaseLevel -- Level of the recursive Id
323 -> [Id] -- Ids that are scrutinised between the binding
324 -- of the recursive Id and here
325 freeScruts env rec_bind_lvl
326 = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
329 %************************************************************************
333 %************************************************************************
336 type LibCaseLevel = Int
338 topLevel :: LibCaseLevel
345 lc_size :: Int, -- Bomb-out size for deciding if
346 -- potential liberatees are too big.
347 -- (passed in from cmd-line args)
349 lc_lvl :: LibCaseLevel, -- Current level
351 lc_lvl_env :: IdEnv LibCaseLevel,
352 -- Binds all non-top-level in-scope Ids
353 -- (top-level and imported things have
356 lc_rec_env :: IdEnv CoreBind,
357 -- Binds *only* recursively defined ids,
358 -- to their own binding group,
359 -- and *only* in their own RHSs
361 lc_scruts :: [(Id,LibCaseLevel)]
362 -- Each of these Ids was scrutinised by an
363 -- enclosing case expression, with the
364 -- specified number of enclosing
365 -- recursive bindings; furthermore,
366 -- the Id is bound at a lower level
367 -- than the case expression. The order is
368 -- insignificant; it's a bag really
371 initEnv :: DynFlags -> LibCaseEnv
373 = LibCaseEnv { lc_size = specThreshold dflags,
375 lc_lvl_env = emptyVarEnv,
376 lc_rec_env = emptyVarEnv,
379 bombOutSize = lc_size