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 )
25 import Name ( localiseName )
27 import Util ( notNull )
28 import Data.IORef ( readIORef )
31 The liberate-case transformation
32 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
33 This module walks over @Core@, and looks for @case@ on free variables.
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.
43 => the inner f is replaced.
50 (note the NEED for shadowing)
59 Better code, because 'a' is free inside the inner letrec, rather
60 than needing projection from v.
62 Other examples we'd like to catch with this kind of transformation
68 We'd like to avoid the redundant pattern match, transforming to
72 last (x:(y:ys)) = last' y ys
75 last' _ (y:ys) = last' y ys
77 (is this necessarily an improvement)
83 drop n (x:xs) = drop (n-1) xs
85 Would like to pass n along unboxed.
87 Note [Scrutinee with cast]
88 ~~~~~~~~~~~~~~~~~~~~~~~~~~
90 f = \ t -> case (v `cast` co) of
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.
97 To think about (Apr 94)
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.
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
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
116 letrec f = let g = ... in ...
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).
124 Note [Indexed data types]
125 ~~~~~~~~~~~~~~~~~~~~~~~~~
127 data family T :: * -> *
131 f x = case x of { DEFAULT -> <body> }
133 We would like to change this to
134 f x = case x `cast` co of { TI p -> <body> }
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
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`.
148 We implement this as part of the liberate-case transformation by
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>
156 The "find the representation type" part is done by looking up in the
157 family-instance environment.
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.
162 %************************************************************************
166 %************************************************************************
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)
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' }) }
184 do_prog env (bind:binds) = bind' : do_prog env' binds
186 (env', bind') = libCaseBind env bind
190 %************************************************************************
194 %************************************************************************
199 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
201 libCaseBind env (NonRec binder rhs)
202 = (addBinders env [binder], NonRec binder (libCase env rhs))
204 libCaseBind env (Rec pairs)
205 = (env_body, Rec pairs')
207 (binders, rhss) = unzip pairs
209 env_body = addBinders env binders
211 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
213 env_rhs = if all rhs_small_enough rhss then extended_env else env
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!
219 extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
220 | (binder, rhs) <- pairs ]
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.
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)))
234 rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
235 lIBERATE_BOMB_SIZE = bombOutSize env
243 libCase :: LibCaseEnv
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
254 libCase env (Lam binder body)
255 = Lam binder (libCase (addBinders env [binder]) body)
257 libCase env (Let bind body)
258 = Let bind' (libCase env_body body)
260 (env_body, bind') = libCaseBind env bind
262 libCase env (Case scrut bndr ty alts)
263 = mkCase env (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
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
270 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
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
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
294 libCaseId :: LibCaseEnv -> Id -> CoreExpr
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)
304 rec_id_level = lookupLevel env v
305 free_scruts = freeScruts env rec_id_level
309 %************************************************************************
313 %************************************************************************
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' }
320 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
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' }
328 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
329 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
331 addScrutedVar :: LibCaseEnv
332 -> Id -- This Id is being scrutinised by a case expression
335 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
336 lc_scruts = scruts }) scrut_var
338 = env { lc_scruts = scruts' }
339 -- Add to scruts iff the scrut_var is being scrutinised at
340 -- a deeper level than its defn
344 scruts' = (scrut_var, lvl) : scruts
345 bind_lvl = case lookupVarEnv lvl_env scrut_var of
349 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
350 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
352 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
354 = case lookupVarEnv (lc_lvl_env env) id of
355 Just lvl -> lc_lvl env
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]
366 %************************************************************************
370 %************************************************************************
373 type LibCaseLevel = Int
375 topLevel :: LibCaseLevel
382 lc_size :: Int, -- Bomb-out size for deciding if
383 -- potential liberatees are too big.
384 -- (passed in from cmd-line args)
386 lc_lvl :: LibCaseLevel, -- Current level
388 lc_lvl_env :: IdEnv LibCaseLevel,
389 -- Binds all non-top-level in-scope Ids
390 -- (top-level and imported things have
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
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
407 lc_fams :: FamInstEnvs
408 -- Instance env for indexed data types
411 initEnv :: DynFlags -> FamInstEnvs -> LibCaseEnv
413 = LibCaseEnv { lc_size = specThreshold dflags,
415 lc_lvl_env = emptyVarEnv,
416 lc_rec_env = emptyVarEnv,
420 bombOutSize = lc_size