2 % (c) The AQUA Project, Glasgow University, 1994
4 \section[LiberateCase]{Unroll recursion to allow evals to be lifted from a loop}
8 #include "HsVersions.h"
10 module LiberateCase ( liberateCase ) where
14 import Id ( localiseId, toplevelishId{-debugging-} )
20 import SimplEnv ( UnfoldingGuidance(..) )
24 This module walks over @Core@, and looks for @case@ on free variables.
26 if there is case on a free on the route to the recursive call,
27 then the recursive call is replaced with an unfolding.
36 => the inner f is replaced.
45 (note the NEED for shadowing)
47 => Run Andr\'e's wonder pass ...
54 Better code, because 'a' is free inside the inner letrec, rather
55 than needing projection from v.
58 To think about (Apr 94)
61 Main worry: duplicating code excessively. At the moment we duplicate
62 the entire binding group once at each recursive call. But there may
63 be a group of recursive calls which share a common set of evaluated
64 free variables, in which case the duplication is a plain waste.
66 Another thing we could consider adding is some unfold-threshold thing,
67 so that we'll only duplicate if the size of the group rhss isn't too
73 The ``level'' of a binder tells how many
74 recursive defns lexically enclose the binding
75 A recursive defn "encloses" its RHS, not its
78 letrec f = let g = ... in ...
83 Here, the level of @f@ is zero, the level of @g@ is one,
84 and the level of @h@ is zero (NB not one).
87 type LibCaseLevel = Int
89 topLevel :: LibCaseLevel
96 Int -- Bomb-out size for deciding if
97 -- potential liberatees are too big.
98 -- (passed in from cmd-line args)
100 LibCaseLevel -- Current level
102 (IdEnv LibCaseLevel) -- Binds all non-top-level in-scope Ids
103 -- (top-level and imported things have
106 (IdEnv PlainCoreBinding)-- Binds *only* recursively defined
107 -- Ids, to their own binding group,
108 -- and *only* in their own RHSs
110 [(Id,LibCaseLevel)] -- Each of these Ids was scrutinised by an
111 -- enclosing case expression, with the
112 -- specified number of enclosing
113 -- recursive bindings; furthermore,
114 -- the Id is bound at a lower level
115 -- than the case expression. The
116 -- order is insignificant; it's a bag
119 initEnv :: Int -> LibCaseEnv
120 initEnv bomb_size = LibCaseEnv bomb_size 0 nullIdEnv nullIdEnv []
122 bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
129 liberateCase :: Int -> [PlainCoreBinding] -> [PlainCoreBinding]
130 liberateCase bomb_size prog
131 = do_prog (initEnv bomb_size) prog
134 do_prog env (bind:binds) = bind' : do_prog env' binds
136 (env', bind') = libCaseBind env bind
143 libCaseBind :: LibCaseEnv -> PlainCoreBinding -> (LibCaseEnv, PlainCoreBinding)
145 libCaseBind env (CoNonRec binder rhs)
146 = (addBinders env [binder], CoNonRec binder (libCase env rhs))
148 libCaseBind env (CoRec pairs)
149 = (env_body, CoRec pairs')
151 (binders, rhss) = unzip pairs
153 env_body = addBinders env binders
155 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
157 env_rhs = if all rhs_small_enough rhss then extended_env else env
159 -- We extend the rec-env by binding each Id to its rhs, first
160 -- processing the rhs with an *un-extended* environment, so
161 -- that the same process doesn't occur for ever!
164 = addRecBinds env [ (localiseId binder, libCase env_body rhs)
165 | (binder, rhs) <- pairs ]
167 -- Why "localiseId" above? Because we're creating a new local
168 -- copy of the original binding. In particular, the original
169 -- binding might have been for a TopLevId, and this copy clearly
170 -- will not be top-level!
172 -- It is enough to change just the binder, because subsequent
173 -- simplification will propagate the right info from the binder.
175 -- Why does it matter? Because the codeGen keeps a separate
176 -- environment for top-level Ids, and it is disastrous for it
177 -- to think that something is top-level when it isn't.
180 = case (calcUnfoldingGuidance True{-sccs OK-} lIBERATE_BOMB_SIZE rhs) of
182 _ -> True -- we didn't BOMB, so it must be OK
184 lIBERATE_BOMB_SIZE = bombOutSize env
192 libCase :: LibCaseEnv
196 libCase env (CoLit lit) = CoLit lit
197 libCase env (CoVar v) = mkCoLetsNoUnboxed (libCaseId env v) (CoVar v)
198 libCase env (CoApp fun arg) = mkCoLetsNoUnboxed (libCaseAtom env arg) (CoApp (libCase env fun) arg)
199 libCase env (CoTyApp fun ty) = CoTyApp (libCase env fun) ty
200 libCase env (CoCon con tys args) = mkCoLetsNoUnboxed (libCaseAtoms env args) (CoCon con tys args)
201 libCase env (CoPrim op tys args) = mkCoLetsNoUnboxed (libCaseAtoms env args) (CoPrim op tys args)
202 libCase env (CoTyLam tyvar body) = CoTyLam tyvar (libCase env body)
203 libCase env (CoSCC cc body) = CoSCC cc (libCase env body)
205 libCase env (CoLam binders body)
206 = CoLam binders (libCase env' body)
208 env' = addBinders env binders
210 libCase env (CoLet bind body)
211 = CoLet bind' (libCase env_body body)
213 (env_body, bind') = libCaseBind env bind
215 libCase env (CoCase scrut alts)
216 = CoCase (libCase env scrut) (libCaseAlts env_alts alts)
218 env_alts = case scrut of
219 CoVar scrut_var -> addScrutedVar env scrut_var
228 libCaseAlts env (CoAlgAlts alts deflt)
229 = CoAlgAlts (map do_alt alts) (libCaseDeflt env deflt)
231 do_alt (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
233 libCaseAlts env (CoPrimAlts alts deflt)
234 = CoPrimAlts (map do_alt alts) (libCaseDeflt env deflt)
236 do_alt (lit,rhs) = (lit, libCase env rhs)
238 libCaseDeflt env CoNoDefault
240 libCaseDeflt env (CoBindDefault binder rhs)
241 = CoBindDefault binder (libCase (addBinders env [binder]) rhs)
247 libCaseAtoms :: LibCaseEnv -> [PlainCoreAtom] -> [PlainCoreBinding]
248 libCaseAtoms env atoms = concat [libCaseAtom env atom | atom <- atoms]
250 libCaseAtom :: LibCaseEnv -> PlainCoreAtom -> [PlainCoreBinding]
251 libCaseAtom env (CoVarAtom arg_id) = libCaseId env arg_id
252 libCaseAtom env (CoLitAtom lit) = []
254 libCaseId :: LibCaseEnv -> Id -> [PlainCoreBinding]
256 | maybeToBool maybe_rec_bind && -- It's a use of a recursive thing
257 there_are_free_scruts -- with free vars scrutinised in RHS
264 maybe_rec_bind :: Maybe PlainCoreBinding -- The binding of the recursive thingy
265 maybe_rec_bind = lookupRecId env v
266 Just the_bind = maybe_rec_bind
268 rec_id_level = lookupLevel env v
270 there_are_free_scruts = freeScruts env rec_id_level
278 addBinders :: LibCaseEnv -> [Id] -> LibCaseEnv
279 addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
280 = LibCaseEnv bomb lvl lvl_env' rec_env scruts
282 lvl_env' = growIdEnvList lvl_env (binders `zip` repeat lvl)
284 addRecBinds :: LibCaseEnv -> [(Id,PlainCoreExpr)] -> LibCaseEnv
285 addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
286 = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
289 lvl_env' = growIdEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
290 rec_env' = growIdEnvList rec_env [(binder, CoRec pairs) | (binder,_) <- pairs]
292 addScrutedVar :: LibCaseEnv
293 -> Id -- This Id is being scrutinised by a case expression
296 addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
298 = LibCaseEnv bomb lvl lvl_env rec_env scruts'
299 -- Add to scruts iff the scrut_var is being scrutinised at
300 -- a deeper level than its defn
304 scruts' = (scrut_var, lvl) : scruts
305 bind_lvl = case lookupIdEnv lvl_env scrut_var of
307 Nothing -> --false: ASSERT(toplevelishId scrut_var)
310 lookupRecId :: LibCaseEnv -> Id -> Maybe PlainCoreBinding
311 lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
313 = lookupIdEnv rec_env id
315 = case (lookupIdEnv rec_env id) of
317 xxx -> --false: ASSERT(toplevelishId id)
321 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
322 lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
323 = case lookupIdEnv lvl_env id of
325 Nothing -> ASSERT(toplevelishId id)
328 freeScruts :: LibCaseEnv
329 -> LibCaseLevel -- Level of the recursive Id
330 -> Bool -- True <=> there is an enclosing case of a variable
331 -- bound outside (ie level <=) the recursive Id.
332 freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
333 = not (null free_scruts)
335 free_scruts = [v | (v,lvl) <- scruts, lvl > rec_bind_lvl]