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 CmdLineOpts ( DynFlags, DynFlag(..), opt_LiberateCaseThreshold )
12 import CoreLint ( showPass, endPass )
14 import CoreUnfold ( couldBeSmallEnoughToInline )
17 import UniqFM ( ufmToList )
21 This module walks over @Core@, and looks for @case@ on free variables.
23 if there is case on a free on the route to the recursive call,
24 then the recursive call is replaced with an unfolding.
33 => the inner f is replaced.
42 (note the NEED for shadowing)
53 Better code, because 'a' is free inside the inner letrec, rather
54 than needing projection from v.
56 Other examples we'd like to catch with this kind of transformation
62 We'd like to avoid the redundant pattern match, transforming to
66 last (x:(y:ys)) = last' y ys
69 last' _ (y:ys) = last' y ys
71 (is this necessarily an improvement)
78 drop n (x:xs) = drop (n-1) xs
80 Would like to pass n along unboxed.
83 To think about (Apr 94)
86 Main worry: duplicating code excessively. At the moment we duplicate
87 the entire binding group once at each recursive call. But there may
88 be a group of recursive calls which share a common set of evaluated
89 free variables, in which case the duplication is a plain waste.
91 Another thing we could consider adding is some unfold-threshold thing,
92 so that we'll only duplicate if the size of the group rhss isn't too
98 The ``level'' of a binder tells how many
99 recursive defns lexically enclose the binding
100 A recursive defn "encloses" its RHS, not its
103 letrec f = let g = ... in ...
108 Here, the level of @f@ is zero, the level of @g@ is one,
109 and the level of @h@ is zero (NB not one).
112 type LibCaseLevel = Int
114 topLevel :: LibCaseLevel
121 Int -- Bomb-out size for deciding if
122 -- potential liberatees are too big.
123 -- (passed in from cmd-line args)
125 LibCaseLevel -- Current level
127 (IdEnv LibCaseLevel) -- Binds all non-top-level in-scope Ids
128 -- (top-level and imported things have
131 (IdEnv CoreBind) -- Binds *only* recursively defined
132 -- Ids, to their own binding group,
133 -- and *only* in their own RHSs
135 [(Id,LibCaseLevel)] -- Each of these Ids was scrutinised by an
136 -- enclosing case expression, with the
137 -- specified number of enclosing
138 -- recursive bindings; furthermore,
139 -- the Id is bound at a lower level
140 -- than the case expression. The
141 -- order is insignificant; it's a bag
144 initEnv :: Int -> LibCaseEnv
145 initEnv bomb_size = LibCaseEnv bomb_size 0 emptyVarEnv emptyVarEnv []
147 pprEnv :: LibCaseEnv -> SDoc
148 pprEnv (LibCaseEnv _ lvl lvl_env _ scruts)
149 = vcat [text "LibCaseEnv" <+> int lvl,
150 fsep (map ppr (ufmToList lvl_env)),
151 fsep (map ppr scruts)]
153 bombOutSize (LibCaseEnv bomb_size _ _ _ _) = bomb_size
160 liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
161 liberateCase dflags binds
163 showPass dflags "Liberate case" ;
164 let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
165 endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
166 {- no specific flag for dumping -}
170 do_prog env (bind:binds) = bind' : do_prog env' binds
172 (env', bind') = libCaseBind env bind
179 libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
181 libCaseBind env (NonRec binder rhs)
182 = (addBinders env [binder], NonRec binder (libCase env rhs))
184 libCaseBind env (Rec pairs)
185 = (env_body, Rec pairs')
187 (binders, rhss) = unzip pairs
189 env_body = addBinders env binders
191 pairs' = [(binder, libCase env_rhs rhs) | (binder,rhs) <- pairs]
193 env_rhs = if all rhs_small_enough rhss then extended_env else env
195 -- We extend the rec-env by binding each Id to its rhs, first
196 -- processing the rhs with an *un-extended* environment, so
197 -- that the same process doesn't occur for ever!
199 extended_env = addRecBinds env [ (binder, libCase env_body rhs)
200 | (binder, rhs) <- pairs ]
202 rhs_small_enough rhs = couldBeSmallEnoughToInline lIBERATE_BOMB_SIZE rhs
203 lIBERATE_BOMB_SIZE = bombOutSize env
211 libCase :: LibCaseEnv
215 libCase env (Var v) = libCaseId env v
216 libCase env (Lit lit) = Lit lit
217 libCase env (Type ty) = Type ty
218 libCase env (App fun arg) = App (libCase env fun) (libCase env arg)
219 libCase env (Note note body) = Note note (libCase env body)
221 libCase env (Lam binder body)
222 = Lam binder (libCase (addBinders env [binder]) body)
224 libCase env (Let bind body)
225 = Let bind' (libCase env_body body)
227 (env_body, bind') = libCaseBind env bind
229 libCase env (Case scrut bndr alts)
230 = Case (libCase env scrut) bndr (map (libCaseAlt env_alts) alts)
232 env_alts = addBinders env_with_scrut [bndr]
233 env_with_scrut = case scrut of
234 Var scrut_var -> addScrutedVar env scrut_var
237 libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
243 libCaseId :: LibCaseEnv -> Id -> CoreExpr
245 | Just the_bind <- lookupRecId env v -- It's a use of a recursive thing
246 , not (null free_scruts) -- with free vars scrutinised in RHS
247 = Let the_bind (Var v)
253 rec_id_level = lookupLevel env v
254 free_scruts = freeScruts env rec_id_level
262 addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
263 addBinders (LibCaseEnv bomb lvl lvl_env rec_env scruts) binders
264 = LibCaseEnv bomb lvl lvl_env' rec_env scruts
266 lvl_env' = extendVarEnvList lvl_env (binders `zip` repeat lvl)
268 addRecBinds :: LibCaseEnv -> [(Id,CoreExpr)] -> LibCaseEnv
269 addRecBinds (LibCaseEnv bomb lvl lvl_env rec_env scruts) pairs
270 = LibCaseEnv bomb lvl' lvl_env' rec_env' scruts
273 lvl_env' = extendVarEnvList lvl_env [(binder,lvl) | (binder,_) <- pairs]
274 rec_env' = extendVarEnvList rec_env [(binder, Rec pairs) | (binder,_) <- pairs]
276 addScrutedVar :: LibCaseEnv
277 -> Id -- This Id is being scrutinised by a case expression
280 addScrutedVar env@(LibCaseEnv bomb lvl lvl_env rec_env scruts) scrut_var
282 = LibCaseEnv bomb lvl lvl_env rec_env scruts'
283 -- Add to scruts iff the scrut_var is being scrutinised at
284 -- a deeper level than its defn
288 scruts' = (scrut_var, lvl) : scruts
289 bind_lvl = case lookupVarEnv lvl_env scrut_var of
293 lookupRecId :: LibCaseEnv -> Id -> Maybe CoreBind
294 lookupRecId (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
295 = lookupVarEnv rec_env id
297 lookupLevel :: LibCaseEnv -> Id -> LibCaseLevel
298 lookupLevel (LibCaseEnv bomb lvl lvl_env rec_env scruts) id
299 = case lookupVarEnv lvl_env id of
303 freeScruts :: LibCaseEnv
304 -> LibCaseLevel -- Level of the recursive Id
305 -> [Id] -- Ids that are scrutinised between the binding
306 -- of the recursive Id and here
307 freeScruts (LibCaseEnv bomb lvl lvl_env rec_env scruts) rec_bind_lvl
308 = [v | (v,scrut_lvl) <- scruts, scrut_lvl > rec_bind_lvl]