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 , [(subst, fam_inst)] <- lookupFamInstEnv (lc_fams env) tycon tys
279 rep_tc = famInstTyCon fam_inst
280 rep_tys = map (substTyVar subst) (tyConTyVars rep_tc)
281 bndr' = setIdType bndr (mkTyConApp rep_tc rep_tys)
282 Just co_tc = tyConFamilyCoercion_maybe rep_tc
283 co = mkTyConApp co_tc rep_tys
284 bind = NonRec bndr (Cast (Var bndr') (mkSymCoercion co))
285 in mkCase env (Cast scrut co) bndr' ty [(DEFAULT,[],Let bind rhs)]
286 mkCase env scrut bndr ty alts
287 = Case scrut bndr ty alts
293 libCaseId :: LibCaseEnv -> Id -> CoreExpr
295 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
296 , notNull free_scruts -- with free vars scrutinised in RHS
297 = Let the_bind (Var v)
303 rec_id_level = lookupLevel env v
304 free_scruts = freeScruts env rec_id_level
308 %************************************************************************
312 %************************************************************************
315 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
316 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
317 = env { lc_lvl_env = lvl_env' }
319 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
321 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
322 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
323 lc_rec_env = rec_env}) pairs
324 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
327 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
328 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
330 addScrutedVar :: LibCaseEnv
331 -> Id -- This Id is being scrutinised by a case expression
334 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
335 lc_scruts = scruts }) scrut_var
337 = env { lc_scruts = scruts' }
338 -- Add to scruts iff the scrut_var is being scrutinised at
339 -- a deeper level than its defn
343 scruts' = (scrut_var, lvl) : scruts
344 bind_lvl = case lookupVarEnv lvl_env scrut_var of
348 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
349 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
351 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
353 = case lookupVarEnv (lc_lvl_env env) id of
354 Just lvl -> lc_lvl env
357 freeScruts :: LibCaseEnv
358 -> LibCaseLevel -- Level of the recursive Id
359 -> [Id] -- Ids that are scrutinised between the binding
360 -- of the recursive Id and here
361 freeScruts env rec_bind_lvl
362 = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
365 %************************************************************************
369 %************************************************************************
372 type LibCaseLevel = Int
374 topLevel :: LibCaseLevel
381 lc_size :: Int, -- Bomb-out size for deciding if
382 -- potential liberatees are too big.
383 -- (passed in from cmd-line args)
385 lc_lvl :: LibCaseLevel, -- Current level
387 lc_lvl_env :: IdEnv LibCaseLevel,
388 -- Binds all non-top-level in-scope Ids
389 -- (top-level and imported things have
392 lc_rec_env :: IdEnv CoreBind,
393 -- Binds *only* recursively defined ids,
394 -- to their own binding group,
395 -- and *only* in their own RHSs
397 lc_scruts :: [(Id,LibCaseLevel)],
398 -- Each of these Ids was scrutinised by an
399 -- enclosing case expression, with the
400 -- specified number of enclosing
401 -- recursive bindings; furthermore,
402 -- the Id is bound at a lower level
403 -- than the case expression. The order is
404 -- insignificant; it's a bag really
406 lc_fams :: FamInstEnvs
407 -- Instance env for indexed data types
410 initEnv :: DynFlags -> FamInstEnvs -> LibCaseEnv
412 = LibCaseEnv { lc_size = specThreshold dflags,
414 lc_lvl_env = emptyVarEnv,
415 lc_rec_env = emptyVarEnv,
419 bombOutSize = lc_size