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 To think about (Apr 94)
94 Main worry: duplicating code excessively. At the moment we duplicate
95 the entire binding group once at each recursive call. But there may
96 be a group of recursive calls which share a common set of evaluated
97 free variables, in which case the duplication is a plain waste.
99 Another thing we could consider adding is some unfold-threshold thing,
100 so that we'll only duplicate if the size of the group rhss isn't too
105 The ``level'' of a binder tells how many
106 recursive defns lexically enclose the binding
107 A recursive defn "encloses" its RHS, not its
110 letrec f = let g = ... in ...
115 Here, the level of @f@ is zero, the level of @g@ is one,
116 and the level of @h@ is zero (NB not one).
119 %************************************************************************
123 %************************************************************************
126 liberateCase :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
127 -> IO (SimplCount, ModGuts)
128 liberateCase hsc_env _ _ guts
129 = do { let dflags = hsc_dflags hsc_env
131 ; showPass dflags "Liberate case"
132 ; let { env = initEnv dflags
133 ; binds' = do_prog env (mg_binds guts) }
134 ; endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
135 {- no specific flag for dumping -}
136 ; return (zeroSimplCount dflags, guts { mg_binds = binds' }) }
139 do_prog env (bind:binds) = bind' : do_prog env' binds
141 (env', bind') = libCaseBind env bind
145 %************************************************************************
149 %************************************************************************
154 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
156 libCaseBind env (NonRec binder rhs)
157 = (addBinders env [binder], NonRec binder (libCase env rhs))
159 libCaseBind env (Rec pairs)
160 = (env_body, Rec pairs')
162 (binders, rhss) = unzip pairs
164 env_body = addBinders env binders
166 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
168 env_rhs = if all rhs_small_enough rhss then extended_env else env
170 -- We extend the rec-env by binding each Id to its rhs, first
171 -- processing the rhs with an *un-extended* environment, so
172 -- that the same process doesn't occur for ever!
174 extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
175 | (binder, rhs) <- pairs ]
177 -- Two subtle things:
178 -- (a) Reset the export flags on the binders so
179 -- that we don't get name clashes on exported things if the
180 -- local binding floats out to top level. This is most unlikely
181 -- to happen, since the whole point concerns free variables.
182 -- But resetting the export flag is right regardless.
184 -- (b) Make the name an Internal one. External Names should never be
185 -- nested; if it were floated to the top level, we'd get a name
186 -- clash at code generation time.
187 adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
189 rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
190 lIBERATE_BOMB_SIZE = bombOutSize env
198 libCase :: LibCaseEnv
202 libCase env (Var v) = libCaseId env v
203 libCase env (Lit lit) = Lit lit
204 libCase env (Type ty) = Type ty
205 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
206 libCase env (Note note body) = Note note (libCase env body)
207 libCase env (Cast e co) = Cast (libCase env e) co
209 libCase env (Lam binder body)
210 = Lam binder (libCase (addBinders env [binder]) body)
212 libCase env (Let bind body)
213 = Let bind' (libCase env_body body)
215 (env_body, bind') = libCaseBind env bind
217 libCase env (Case scrut bndr ty alts)
218 = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
220 env_alts = addBinders (mk_alt_env scrut) [bndr]
221 mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
222 mk_alt_env (Cast scrut _) = mk_alt_env scrut -- Note [Scrutinee with cast]
223 mk_alt_env otehr = env
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
247 %************************************************************************
251 %************************************************************************
254 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
255 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
256 = env { lc_lvl_env = lvl_env' }
258 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
260 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
261 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
262 lc_rec_env = rec_env}) pairs
263 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
266 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
267 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
269 addScrutedVar :: LibCaseEnv
270 -> Id -- This Id is being scrutinised by a case expression
273 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
274 lc_scruts = scruts }) scrut_var
276 = env { lc_scruts = scruts' }
277 -- Add to scruts iff the scrut_var is being scrutinised at
278 -- a deeper level than its defn
282 scruts' = (scrut_var, lvl) : scruts
283 bind_lvl = case lookupVarEnv lvl_env scrut_var of
287 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
288 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
290 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
292 = case lookupVarEnv (lc_lvl_env env) id of
296 freeScruts :: LibCaseEnv
297 -> LibCaseLevel -- Level of the recursive Id
298 -> [Id] -- Ids that are scrutinised between the binding
299 -- of the recursive Id and here
300 freeScruts env rec_bind_lvl
301 = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
304 %************************************************************************
308 %************************************************************************
311 type LibCaseLevel = Int
313 topLevel :: LibCaseLevel
320 lc_size :: Int, -- Bomb-out size for deciding if
321 -- potential liberatees are too big.
322 -- (passed in from cmd-line args)
324 lc_lvl :: LibCaseLevel, -- Current level
326 lc_lvl_env :: IdEnv LibCaseLevel,
327 -- Binds all non-top-level in-scope Ids
328 -- (top-level and imported things have
331 lc_rec_env :: IdEnv CoreBind,
332 -- Binds *only* recursively defined ids,
333 -- to their own binding group,
334 -- and *only* in their own RHSs
336 lc_scruts :: [(Id,LibCaseLevel)]
337 -- Each of these Ids was scrutinised by an
338 -- enclosing case expression, with the
339 -- specified number of enclosing
340 -- recursive bindings; furthermore,
341 -- the Id is bound at a lower level
342 -- than the case expression. The order is
343 -- insignificant; it's a bag really
346 initEnv :: DynFlags -> LibCaseEnv
348 = LibCaseEnv { lc_size = specThreshold dflags,
350 lc_lvl_env = emptyVarEnv,
351 lc_rec_env = emptyVarEnv,
354 bombOutSize = lc_size