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 optimistaion (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
112 The ``level'' of a binder tells how many
113 recursive defns lexically enclose the binding
114 A recursive defn "encloses" its RHS, not its
117 letrec f = let g = ... in ...
122 Here, the level of @f@ is zero, the level of @g@ is one,
123 and the level of @h@ is zero (NB not one).
125 Note [Indexed data types]
126 ~~~~~~~~~~~~~~~~~~~~~~~~~
128 data family T :: * -> *
132 f x = case x of { DEFAULT -> <body> }
134 We would like to change this to
135 f x = case x `cast` co of { TI p -> <body> }
137 so that <body> can make use of the fact that x is already evaluated to
138 a TI; and a case on a known data type may be more efficient than a
139 polymorphic one (not sure this is true any longer). Anyway the former
140 showed up in Roman's experiments. Example:
141 foo :: FooT Int -> Int -> Int
142 foo t n = t `seq` bar n
145 bar n = bar (n - case t of TI i -> i)
146 Here we'd like to avoid repeated evaluating t inside the loop, by
147 taking advantage of the `seq`.
149 We implement this as part of the liberate-case transformation by
151 case <scrut> of (x::T) tys { DEFAULT -> <body> }
152 where x :: T tys, and T is a indexed family tycon. Find the
153 representation type (T77 tys'), and coercion co, and transform to
154 case <scrut> `cast` co of (y::T77 tys')
155 DEFAULT -> let x = y `cast` sym co in <body>
157 The "find the representation type" part is done by looking up in the
158 family-instance environment.
160 NB: in fact we re-use x (changing its type) to avoid making a fresh y;
161 this entails shadowing, but that's ok.
163 %************************************************************************
167 %************************************************************************
170 liberateCase :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
171 -> IO (SimplCount, ModGuts)
172 liberateCase hsc_env _ _ guts
173 = do { let dflags = hsc_dflags hsc_env
174 ; eps <- readIORef (hsc_EPS hsc_env)
175 ; let fam_envs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
177 ; showPass dflags "Liberate case"
178 ; let { env = initEnv dflags fam_envs
179 ; binds' = do_prog env (mg_binds guts) }
180 ; endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
181 {- no specific flag for dumping -}
182 ; return (zeroSimplCount dflags, guts { mg_binds = binds' }) }
185 do_prog env (bind:binds) = bind' : do_prog env' binds
187 (env', bind') = libCaseBind env bind
191 %************************************************************************
195 %************************************************************************
200 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
202 libCaseBind env (NonRec binder rhs)
203 = (addBinders env [binder], NonRec binder (libCase env rhs))
205 libCaseBind env (Rec pairs)
206 = (env_body, Rec pairs')
208 (binders, rhss) = unzip pairs
210 env_body = addBinders env binders
212 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
214 env_rhs = if all rhs_small_enough rhss then extended_env else env
216 -- We extend the rec-env by binding each Id to its rhs, first
217 -- processing the rhs with an *un-extended* environment, so
218 -- that the same process doesn't occur for ever!
220 extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
221 | (binder, rhs) <- pairs ]
223 -- Two subtle things:
224 -- (a) Reset the export flags on the binders so
225 -- that we don't get name clashes on exported things if the
226 -- local binding floats out to top level. This is most unlikely
227 -- to happen, since the whole point concerns free variables.
228 -- But resetting the export flag is right regardless.
230 -- (b) Make the name an Internal one. External Names should never be
231 -- nested; if it were floated to the top level, we'd get a name
232 -- clash at code generation time.
233 adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
235 rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
236 lIBERATE_BOMB_SIZE = bombOutSize env
244 libCase :: LibCaseEnv
248 libCase env (Var v) = libCaseId env v
249 libCase env (Lit lit) = Lit lit
250 libCase env (Type ty) = Type ty
251 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
252 libCase env (Note note body) = Note note (libCase env body)
253 libCase env (Cast e co) = Cast (libCase env e) co
255 libCase env (Lam binder body)
256 = Lam binder (libCase (addBinders env [binder]) body)
258 libCase env (Let bind body)
259 = Let bind' (libCase env_body body)
261 (env_body, bind') = libCaseBind env bind
263 libCase env (Case scrut bndr ty alts)
264 = mkCase env (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
266 env_alts = addBinders (mk_alt_env scrut) [bndr]
267 mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
268 mk_alt_env (Cast scrut _) = mk_alt_env scrut -- Note [Scrutinee with cast]
269 mk_alt_env otehr = env
271 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
275 mkCase :: LibCaseEnv -> CoreExpr -> Id -> Type -> [CoreAlt] -> CoreExpr
276 -- See Note [Indexed data types]
277 mkCase env scrut bndr ty [(DEFAULT,_,rhs)]
278 | Just (tycon, tys) <- splitTyConApp_maybe (idType bndr)
279 , [(subst, fam_inst)] <- lookupFamInstEnv (lc_fams env) tycon tys
281 rep_tc = famInstTyCon fam_inst
282 rep_tys = map (substTyVar subst) (tyConTyVars rep_tc)
283 bndr' = setIdType bndr (mkTyConApp rep_tc rep_tys)
284 Just co_tc = tyConFamilyCoercion_maybe rep_tc
285 co = mkTyConApp co_tc rep_tys
286 bind = NonRec bndr (Cast (Var bndr') (mkSymCoercion co))
287 in mkCase env (Cast scrut co) bndr' ty [(DEFAULT,[],Let bind rhs)]
288 mkCase env scrut bndr ty alts
289 = Case scrut bndr ty alts
295 libCaseId :: LibCaseEnv -> Id -> CoreExpr
297 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
298 , notNull free_scruts -- with free vars scrutinised in RHS
299 = Let the_bind (Var v)
305 rec_id_level = lookupLevel env v
306 free_scruts = freeScruts env rec_id_level
310 %************************************************************************
314 %************************************************************************
317 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
318 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
319 = env { lc_lvl_env = lvl_env' }
321 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
323 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
324 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
325 lc_rec_env = rec_env}) pairs
326 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
329 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
330 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
332 addScrutedVar :: LibCaseEnv
333 -> Id -- This Id is being scrutinised by a case expression
336 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
337 lc_scruts = scruts }) scrut_var
339 = env { lc_scruts = scruts' }
340 -- Add to scruts iff the scrut_var is being scrutinised at
341 -- a deeper level than its defn
345 scruts' = (scrut_var, lvl) : scruts
346 bind_lvl = case lookupVarEnv lvl_env scrut_var of
350 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
351 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
353 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
355 = case lookupVarEnv (lc_lvl_env env) id of
356 Just lvl -> lc_lvl env
359 freeScruts :: LibCaseEnv
360 -> LibCaseLevel -- Level of the recursive Id
361 -> [Id] -- Ids that are scrutinised between the binding
362 -- of the recursive Id and here
363 freeScruts env rec_bind_lvl
364 = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
367 %************************************************************************
371 %************************************************************************
374 type LibCaseLevel = Int
376 topLevel :: LibCaseLevel
383 lc_size :: Int, -- Bomb-out size for deciding if
384 -- potential liberatees are too big.
385 -- (passed in from cmd-line args)
387 lc_lvl :: LibCaseLevel, -- Current level
389 lc_lvl_env :: IdEnv LibCaseLevel,
390 -- Binds all non-top-level in-scope Ids
391 -- (top-level and imported things have
394 lc_rec_env :: IdEnv CoreBind,
395 -- Binds *only* recursively defined ids,
396 -- to their own binding group,
397 -- and *only* in their own RHSs
399 lc_scruts :: [(Id,LibCaseLevel)],
400 -- Each of these Ids was scrutinised by an
401 -- enclosing case expression, with the
402 -- specified number of enclosing
403 -- recursive bindings; furthermore,
404 -- the Id is bound at a lower level
405 -- than the case expression. The order is
406 -- insignificant; it's a bag really
408 lc_fams :: FamInstEnvs
409 -- Instance env for indexed data types
412 initEnv :: DynFlags -> FamInstEnvs -> LibCaseEnv
414 = LibCaseEnv { lc_size = libCaseThreshold dflags,
416 lc_lvl_env = emptyVarEnv,
417 lc_rec_env = emptyVarEnv,
421 bombOutSize = lc_size