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"
11 import DynFlags ( DynFlags, DynFlag(..) )
12 import StaticFlags ( opt_LiberateCaseThreshold )
13 import CoreLint ( showPass, endPass )
15 import CoreUnfold ( couldBeSmallEnoughToInline )
16 import Id ( Id, setIdName, idName, setIdNotExported )
18 import Name ( localiseName )
20 import Util ( notNull )
23 This module walks over @Core@, and looks for @case@ on free variables.
25 if there is case on a free on the route to the recursive call,
26 then the recursive call is replaced with an unfolding.
35 => the inner f is replaced.
44 (note the NEED for shadowing)
55 Better code, because 'a' is free inside the inner letrec, rather
56 than needing projection from v.
58 Other examples we'd like to catch with this kind of transformation
64 We'd like to avoid the redundant pattern match, transforming to
68 last (x:(y:ys)) = last' y ys
71 last' _ (y:ys) = last' y ys
73 (is this necessarily an improvement)
80 drop n (x:xs) = drop (n-1) xs
82 Would like to pass n along unboxed.
84 Note [Scrutinee with cast]
85 ~~~~~~~~~~~~~~~~~~~~~~~~~~
87 f = \ t -> case (v `cast` co) of
90 Exactly the same optimistaion (unrolling one call to f) will work here,
91 despite the cast. See mk_alt_env in the Case branch of libCase.
94 To think about (Apr 94)
97 Main worry: duplicating code excessively. At the moment we duplicate
98 the entire binding group once at each recursive call. But there may
99 be a group of recursive calls which share a common set of evaluated
100 free variables, in which case the duplication is a plain waste.
102 Another thing we could consider adding is some unfold-threshold thing,
103 so that we'll only duplicate if the size of the group rhss isn't too
109 The ``level'' of a binder tells how many
110 recursive defns lexically enclose the binding
111 A recursive defn "encloses" its RHS, not its
114 letrec f = let g = ... in ...
119 Here, the level of @f@ is zero, the level of @g@ is one,
120 and the level of @h@ is zero (NB not one).
123 type LibCaseLevel = Int
125 topLevel :: LibCaseLevel
132 lc_size :: Int, -- Bomb-out size for deciding if
133 -- potential liberatees are too big.
134 -- (passed in from cmd-line args)
136 lc_lvl :: LibCaseLevel, -- Current level
138 lc_lvl_env :: IdEnv LibCaseLevel,
139 -- Binds all non-top-level in-scope Ids
140 -- (top-level and imported things have
143 lc_rec_env :: IdEnv CoreBind,
144 -- Binds *only* recursively defined ids,
145 -- to their own binding group,
146 -- and *only* in their own RHSs
148 lc_scruts :: [(Id,LibCaseLevel)]
149 -- Each of these Ids was scrutinised by an
150 -- enclosing case expression, with the
151 -- specified number of enclosing
152 -- recursive bindings; furthermore,
153 -- the Id is bound at a lower level
154 -- than the case expression. The order is
155 -- insignificant; it's a bag really
157 -- lc_fams :: FamInstEnvs
158 -- Instance env for indexed data types
161 initEnv :: Int -> LibCaseEnv
163 = LibCaseEnv { lc_size = bomb_size, lc_lvl = 0,
164 lc_lvl_env = emptyVarEnv, lc_rec_env = emptyVarEnv,
167 bombOutSize = lc_size
174 liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
175 liberateCase dflags binds
177 showPass dflags "Liberate case" ;
178 let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
179 endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
180 {- no specific flag for dumping -}
184 do_prog env (bind:binds) = bind' : do_prog env' binds
186 (env', bind') = libCaseBind env bind
193 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
195 libCaseBind env (NonRec binder rhs)
196 = (addBinders env [binder], NonRec binder (libCase env rhs))
198 libCaseBind env (Rec pairs)
199 = (env_body, Rec pairs')
201 (binders, rhss) = unzip pairs
203 env_body = addBinders env binders
205 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
207 env_rhs = if all rhs_small_enough rhss then extended_env else env
209 -- We extend the rec-env by binding each Id to its rhs, first
210 -- processing the rhs with an *un-extended* environment, so
211 -- that the same process doesn't occur for ever!
213 extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
214 | (binder, rhs) <- pairs ]
216 -- Two subtle things:
217 -- (a) Reset the export flags on the binders so
218 -- that we don't get name clashes on exported things if the
219 -- local binding floats out to top level. This is most unlikely
220 -- to happen, since the whole point concerns free variables.
221 -- But resetting the export flag is right regardless.
223 -- (b) Make the name an Internal one. External Names should never be
224 -- nested; if it were floated to the top level, we'd get a name
225 -- clash at code generation time.
226 adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
228 rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
229 lIBERATE_BOMB_SIZE = bombOutSize env
237 libCase :: LibCaseEnv
241 libCase env (Var v) = libCaseId env v
242 libCase env (Lit lit) = Lit lit
243 libCase env (Type ty) = Type ty
244 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
245 libCase env (Note note body) = Note note (libCase env body)
246 libCase env (Cast e co) = Cast (libCase env e) co
248 libCase env (Lam binder body)
249 = Lam binder (libCase (addBinders env [binder]) body)
251 libCase env (Let bind body)
252 = Let bind' (libCase env_body body)
254 (env_body, bind') = libCaseBind env bind
256 libCase env (Case scrut bndr ty alts)
257 = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
259 env_alts = addBinders (mk_alt_env scrut) [bndr]
260 mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
261 mk_alt_env (Cast scrut _) = mk_alt_env scrut -- Note [Scrutinee with cast]
262 mk_alt_env otehr = env
264 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
270 libCaseId :: LibCaseEnv -> Id -> CoreExpr
272 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
273 , notNull free_scruts -- with free vars scrutinised in RHS
274 = Let the_bind (Var v)
280 rec_id_level = lookupLevel env v
281 free_scruts = freeScruts env rec_id_level
289 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
290 addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
291 = env { lc_lvl_env = lvl_env' }
293 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
295 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
296 addRecBinds env@(LibCaseEnv {lc_lvl = lvl, lc_lvl_env = lvl_env,
297 lc_rec_env = rec_env}) pairs
298 = env { lc_lvl = lvl', lc_lvl_env = lvl_env', lc_rec_env = rec_env' }
301 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
302 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
304 addScrutedVar :: LibCaseEnv
305 -> Id -- This Id is being scrutinised by a case expression
308 addScrutedVar env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env,
309 lc_scruts = scruts }) scrut_var
311 = env { lc_scruts = scruts' }
312 -- Add to scruts iff the scrut_var is being scrutinised at
313 -- a deeper level than its defn
317 scruts' = (scrut_var, lvl) : scruts
318 bind_lvl = case lookupVarEnv lvl_env scrut_var of
322 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
323 lookupRecId env id = lookupVarEnv (lc_rec_env env) id
325 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
327 = case lookupVarEnv (lc_lvl_env env) id of
328 Just lvl -> lc_lvl env
331 freeScruts :: LibCaseEnv
332 -> LibCaseLevel -- Level of the recursive Id
333 -> [Id] -- Ids that are scrutinised between the binding
334 -- of the recursive Id and here
335 freeScruts env rec_bind_lvl
336 = [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]