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 )
26 import Util ( notNull )
27 import Data.IORef ( readIORef )
30 The liberate-case transformation
31 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
32 This module walks over @Core@, and looks for @case@ on free variables.
34 if there is case on a free on the route to the recursive call,
35 then the recursive call is replaced with an unfolding.
42 => the inner f is replaced.
49 (note the NEED for shadowing)
58 Better code, because 'a' is free inside the inner letrec, rather
59 than needing projection from v.
61 Other examples we'd like to catch with this kind of transformation
67 We'd like to avoid the redundant pattern match, transforming to
71 last (x:(y:ys)) = last' y ys
74 last' _ (y:ys) = last' y ys
76 (is this necessarily an improvement)
82 drop n (x:xs) = drop (n-1) xs
84 Would like to pass n along unboxed.
86 Note [Scrutinee with cast]
87 ~~~~~~~~~~~~~~~~~~~~~~~~~~
89 f = \ t -> case (v `cast` co) of
92 Exactly the same optimisation (unrolling one call to f) will work here,
93 despite the cast. See mk_alt_env in the Case branch of libCase.
96 To think about (Apr 94)
99 Main worry: duplicating code excessively. At the moment we duplicate
100 the entire binding group once at each recursive call. But there may
101 be a group of recursive calls which share a common set of evaluated
102 free variables, in which case the duplication is a plain waste.
104 Another thing we could consider adding is some unfold-threshold thing,
105 so that we'll only duplicate if the size of the group rhss isn't too
110 The ``level'' of a binder tells how many
111 recursive defns lexically enclose the binding
112 A recursive defn "encloses" its RHS, not its
115 letrec f = let g = ... in ...
120 Here, the level of @f@ is zero, the level of @g@ is one,
121 and the level of @h@ is zero (NB not one).
123 Note [Indexed data types]
124 ~~~~~~~~~~~~~~~~~~~~~~~~~
126 data family T :: * -> *
130 f x = case x of { DEFAULT -> <body> }
132 We would like to change this to
133 f x = case x `cast` co of { TI p -> <body> }
135 so that <body> can make use of the fact that x is already evaluated to
136 a TI; and a case on a known data type may be more efficient than a
137 polymorphic one (not sure this is true any longer). Anyway the former
138 showed up in Roman's experiments. Example:
139 foo :: FooT Int -> Int -> Int
140 foo t n = t `seq` bar n
143 bar n = bar (n - case t of TI i -> i)
144 Here we'd like to avoid repeated evaluating t inside the loop, by
145 taking advantage of the `seq`.
147 We implement this as part of the liberate-case transformation by
149 case <scrut> of (x::T) tys { DEFAULT -> <body> }
150 where x :: T tys, and T is a indexed family tycon. Find the
151 representation type (T77 tys'), and coercion co, and transform to
152 case <scrut> `cast` co of (y::T77 tys')
153 DEFAULT -> let x = y `cast` sym co in <body>
155 The "find the representation type" part is done by looking up in the
156 family-instance environment.
158 NB: in fact we re-use x (changing its type) to avoid making a fresh y;
159 this entails shadowing, but that's ok.
161 %************************************************************************
165 %************************************************************************
168 liberateCase :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
169 -> IO (SimplCount, ModGuts)
170 liberateCase hsc_env _ _ guts
171 = do { let dflags = hsc_dflags hsc_env
172 ; eps <- readIORef (hsc_EPS hsc_env)
173 ; let fam_envs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
175 ; showPass dflags "Liberate case"
176 ; let { env = initEnv dflags fam_envs
177 ; binds' = do_prog env (mg_binds guts) }
178 ; endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
179 {- no specific flag for dumping -}
180 ; return (zeroSimplCount dflags, guts { mg_binds = binds' }) }
183 do_prog env (bind:binds) = bind' : do_prog env' binds
185 (env', bind') = libCaseBind env bind
189 %************************************************************************
193 %************************************************************************
198 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
200 libCaseBind env (NonRec binder rhs)
201 = (addBinders env [binder], NonRec binder (libCase env rhs))
203 libCaseBind env (Rec pairs)
204 = (env_body, Rec pairs')
206 (binders, rhss) = unzip pairs
208 env_body = addBinders env binders
210 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
212 env_rhs = if all rhs_small_enough rhss then extended_env else env
214 -- We extend the rec-env by binding each Id to its rhs, first
215 -- processing the rhs with an *un-extended* environment, so
216 -- that the same process doesn't occur for ever!
218 extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
219 | (binder, rhs) <- pairs ]
221 -- Two subtle things:
222 -- (a) Reset the export flags on the binders so
223 -- that we don't get name clashes on exported things if the
224 -- local binding floats out to top level. This is most unlikely
225 -- to happen, since the whole point concerns free variables.
226 -- But resetting the export flag is right regardless.
228 -- (b) Make the name an Internal one. External Names should never be
229 -- nested; if it were floated to the top level, we'd get a name
230 -- clash at code generation time.
231 adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
233 rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
234 lIBERATE_BOMB_SIZE = bombOutSize env
242 libCase :: LibCaseEnv
246 libCase env (Var v) = libCaseId env v
247 libCase env (Lit lit) = Lit lit
248 libCase env (Type ty) = Type ty
249 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
250 libCase env (Note note body) = Note note (libCase env body)
251 libCase env (Cast e co) = Cast (libCase env e) co
253 libCase env (Lam binder body)
254 = Lam binder (libCase (addBinders env [binder]) body)
256 libCase env (Let bind body)
257 = Let bind' (libCase env_body body)
259 (env_body, bind') = libCaseBind env bind
261 libCase env (Case scrut bndr ty alts)
262 = mkCase env (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
264 env_alts = addBinders (mk_alt_env scrut) [bndr]
265 mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
266 mk_alt_env (Cast scrut _) = mk_alt_env scrut -- Note [Scrutinee with cast]
267 mk_alt_env otehr = env
269 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
273 mkCase :: LibCaseEnv -> CoreExpr -> Id -> Type -> [CoreAlt] -> CoreExpr
274 -- See Note [Indexed data types]
275 mkCase env scrut bndr ty [(DEFAULT,_,rhs)]
276 | Just (tycon, tys) <- splitTyConApp_maybe (idType bndr)
277 , [(fam_inst, rep_tys)] <- lookupFamInstEnv (lc_fams env) tycon tys
279 rep_tc = famInstTyCon fam_inst
280 bndr' = setIdType bndr (mkTyConApp rep_tc rep_tys)
281 Just co_tc = tyConFamilyCoercion_maybe rep_tc
282 co = mkTyConApp co_tc rep_tys
283 bind = NonRec bndr (Cast (Var bndr') (mkSymCoercion co))
284 in mkCase env (Cast scrut co) bndr' ty [(DEFAULT,[],Let bind rhs)]
285 mkCase env scrut bndr ty alts
286 = Case scrut bndr ty alts
292 libCaseId :: LibCaseEnv -> Id -> CoreExpr
294 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
295 , notNull free_scruts -- with free vars scrutinised in RHS
296 = Let the_bind (Var v)
302 rec_id_level = lookupLevel env v
303 free_scruts = freeScruts env rec_id_level
307 %************************************************************************
311 %************************************************************************
314 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
315 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
316 = env { lc_lvl_env = lvl_env' }
318 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
320 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
321 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
322 lc_rec_env = rec_env}) pairs
323 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
326 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
327 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
329 addScrutedVar :: LibCaseEnv
330 -> Id -- This Id is being scrutinised by a case expression
333 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
334 lc_scruts = scruts }) scrut_var
336 = env { lc_scruts = scruts' }
337 -- Add to scruts iff the scrut_var is being scrutinised at
338 -- a deeper level than its defn
342 scruts' = (scrut_var, lvl) : scruts
343 bind_lvl = case lookupVarEnv lvl_env scrut_var of
347 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
348 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
350 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
352 = case lookupVarEnv (lc_lvl_env env) id of
356 freeScruts :: LibCaseEnv
357 -> LibCaseLevel -- Level of the recursive Id
358 -> [Id] -- Ids that are scrutinised between the binding
359 -- of the recursive Id and here
360 freeScruts env rec_bind_lvl
361 = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
364 %************************************************************************
368 %************************************************************************
371 type LibCaseLevel = Int
373 topLevel :: LibCaseLevel
380 lc_size :: Int, -- Bomb-out size for deciding if
381 -- potential liberatees are too big.
382 -- (passed in from cmd-line args)
384 lc_lvl :: LibCaseLevel, -- Current level
386 lc_lvl_env :: IdEnv LibCaseLevel,
387 -- Binds all non-top-level in-scope Ids
388 -- (top-level and imported things have
391 lc_rec_env :: IdEnv CoreBind,
392 -- Binds *only* recursively defined ids,
393 -- to their own binding group,
394 -- and *only* in their own RHSs
396 lc_scruts :: [(Id,LibCaseLevel)],
397 -- Each of these Ids was scrutinised by an
398 -- enclosing case expression, with the
399 -- specified number of enclosing
400 -- recursive bindings; furthermore,
401 -- the Id is bound at a lower level
402 -- than the case expression. The order is
403 -- insignificant; it's a bag really
405 lc_fams :: FamInstEnvs
406 -- Instance env for indexed data types
409 initEnv :: DynFlags -> FamInstEnvs -> LibCaseEnv
411 = LibCaseEnv { lc_size = specThreshold dflags,
413 lc_lvl_env = emptyVarEnv,
414 lc_rec_env = emptyVarEnv,
418 bombOutSize = lc_size