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.
85 To think about (Apr 94)
88 Main worry: duplicating code excessively. At the moment we duplicate
89 the entire binding group once at each recursive call. But there may
90 be a group of recursive calls which share a common set of evaluated
91 free variables, in which case the duplication is a plain waste.
93 Another thing we could consider adding is some unfold-threshold thing,
94 so that we'll only duplicate if the size of the group rhss isn't too
100 The ``level'' of a binder tells how many
101 recursive defns lexically enclose the binding
102 A recursive defn "encloses" its RHS, not its
105 letrec f = let g = ... in ...
110 Here, the level of @f@ is zero, the level of @g@ is one,
111 and the level of @h@ is zero (NB not one).
114 type LibCaseLevel = Int
116 topLevel :: LibCaseLevel
123 Int -- Bomb-out size for deciding if
124 -- potential liberatees are too big.
125 -- (passed in from cmd-line args)
127 LibCaseLevel -- Current level
129 (IdEnv LibCaseLevel) -- Binds all non-top-level in-scope Ids
130 -- (top-level and imported things have
133 (IdEnv CoreBind) -- Binds *only* recursively defined
134 -- Ids, to their own binding group,
135 -- and *only* in their own RHSs
137 [(Id,LibCaseLevel)] -- Each of these Ids was scrutinised by an
138 -- enclosing case expression, with the
139 -- specified number of enclosing
140 -- recursive bindings; furthermore,
141 -- the Id is bound at a lower level
142 -- than the case expression. The
143 -- order is insignificant; it's a bag
146 initEnv :: Int -> LibCaseEnv
147 initEnv bomb_size = LibCaseEnv bomb_size 0 emptyVarEnv emptyVarEnv []
149 bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
156 liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
157 liberateCase dflags binds
159 showPass dflags "Liberate case" ;
160 let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
161 endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
162 {- no specific flag for dumping -}
166 do_prog env (bind:binds) = bind' : do_prog env' binds
168 (env', bind') = libCaseBind env bind
175 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
177 libCaseBind env (NonRec binder rhs)
178 = (addBinders env [binder], NonRec binder (libCase env rhs))
180 libCaseBind env (Rec pairs)
181 = (env_body, Rec pairs')
183 (binders, rhss) = unzip pairs
185 env_body = addBinders env binders
187 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
189 env_rhs = if all rhs_small_enough rhss then extended_env else env
191 -- We extend the rec-env by binding each Id to its rhs, first
192 -- processing the rhs with an *un-extended* environment, so
193 -- that the same process doesn't occur for ever!
195 extended_env = addRecBinds env [ (adjust binder, libCase env_body rhs)
196 | (binder, rhs) <- pairs ]
198 -- Two subtle things:
199 -- (a) Reset the export flags on the binders so
200 -- that we don't get name clashes on exported things if the
201 -- local binding floats out to top level. This is most unlikely
202 -- to happen, since the whole point concerns free variables.
203 -- But resetting the export flag is right regardless.
205 -- (b) Make the name an Internal one. External Names should never be
206 -- nested; if it were floated to the top level, we'd get a name
207 -- clash at code generation time.
208 adjust bndr = setIdNotExported (setIdName bndr (localiseName (idName bndr)))
210 rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
211 lIBERATE_BOMB_SIZE = bombOutSize env
219 libCase :: LibCaseEnv
223 libCase env (Var v) = libCaseId env v
224 libCase env (Lit lit) = Lit lit
225 libCase env (Type ty) = Type ty
226 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
227 libCase env (Note note body) = Note note (libCase env body)
228 libCase env (Cast e co) = Cast (libCase env e) co
230 libCase env (Lam binder body)
231 = Lam binder (libCase (addBinders env [binder]) body)
233 libCase env (Let bind body)
234 = Let bind' (libCase env_body body)
236 (env_body, bind') = libCaseBind env bind
238 libCase env (Case scrut bndr ty alts)
239 = Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
241 env_alts = addBinders env_with_scrut [bndr]
242 env_with_scrut = case scrut of
243 Var scrut_var -> addScrutedVar env scrut_var
246 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
252 libCaseId :: LibCaseEnv -> Id -> CoreExpr
254 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
255 , notNull free_scruts -- with free vars scrutinised in RHS
256 = Let the_bind (Var v)
262 rec_id_level = lookupLevel env v
263 free_scruts = freeScruts env rec_id_level
271 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
272 addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
273 = LibCaseEnv bomb lvl lvl_env' rec_env scruts
275 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
277 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
278 addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
279 = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
282 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
283 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
285 addScrutedVar :: LibCaseEnv
286 -> Id -- This Id is being scrutinised by a case expression
289 addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
291 = LibCaseEnv bomb lvl lvl_env rec_env scruts'
292 -- Add to scruts iff the scrut_var is being scrutinised at
293 -- a deeper level than its defn
297 scruts' = (scrut_var, lvl) : scruts
298 bind_lvl = case lookupVarEnv lvl_env scrut_var of
302 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
303 lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
304 = lookupVarEnv rec_env id
306 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
307 lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
308 = case lookupVarEnv lvl_env id of
312 freeScruts :: LibCaseEnv
313 -> LibCaseLevel -- Level of the recursive Id
314 -> [Id] -- Ids that are scrutinised between the binding
315 -- of the recursive Id and here
316 freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
317 = [v | (v,scrut_lvl) <- scruts, scrut_lvl > rec_bind_lvl]