Join points
[ghc.git] / compiler / simplCore / SimplEnv.hs
1 {-
2 (c) The AQUA Project, Glasgow University, 1993-1998
3
4 \section[SimplMonad]{The simplifier Monad}
5 -}
6
7 {-# LANGUAGE CPP #-}
8
9 module SimplEnv (
10 -- * The simplifier mode
11 setMode, getMode, updMode,
12
13 -- * Environments
14 SimplEnv(..), StaticEnv, pprSimplEnv, -- Temp not abstract
15 mkSimplEnv, extendIdSubst,
16 SimplEnv.extendTvSubst, SimplEnv.extendCvSubst,
17 zapSubstEnv, setSubstEnv,
18 getInScope, setInScope, setInScopeSet, modifyInScope, addNewInScopeIds,
19 getSimplRules,
20
21 -- * Substitution results
22 SimplSR(..), mkContEx, substId, lookupRecBndr, refineFromInScope,
23 isJoinIdInEnv_maybe,
24
25 -- * Simplifying 'Id' binders
26 simplNonRecBndr, simplNonRecJoinBndr, simplRecBndrs, simplRecJoinBndrs,
27 simplBinder, simplBinders,
28 substTy, substTyVar, getTCvSubst,
29 substCo, substCoVar,
30
31 -- * Floats
32 Floats, emptyFloats, isEmptyFloats,
33 addNonRec, addFloats, extendFloats,
34 wrapFloats, setFloats, zapFloats, addRecFloats, mapFloats,
35 doFloatFromRhs, getFloatBinds,
36
37 JoinFloats, emptyJoinFloats, isEmptyJoinFloats,
38 wrapJoinFloats, zapJoinFloats, restoreJoinFloats, getJoinFloatBinds,
39 ) where
40
41 #include "HsVersions.h"
42
43 import SimplMonad
44 import CoreMonad ( SimplifierMode(..) )
45 import CoreSyn
46 import CoreUtils
47 import Var
48 import VarEnv
49 import VarSet
50 import OrdList
51 import Id
52 import MkCore ( mkWildValBinder )
53 import TysWiredIn
54 import qualified Type
55 import Type hiding ( substTy, substTyVar, substTyVarBndr )
56 import qualified Coercion
57 import Coercion hiding ( substCo, substCoVar, substCoVarBndr )
58 import BasicTypes
59 import MonadUtils
60 import Outputable
61 import Util
62 import UniqFM ( pprUniqFM )
63
64 import Data.List
65
66 {-
67 ************************************************************************
68 * *
69 \subsubsection{The @SimplEnv@ type}
70 * *
71 ************************************************************************
72 -}
73
74 data SimplEnv
75 = SimplEnv {
76 ----------- Static part of the environment -----------
77 -- Static in the sense of lexically scoped,
78 -- wrt the original expression
79
80 seMode :: SimplifierMode,
81
82 -- The current substitution
83 seTvSubst :: TvSubstEnv, -- InTyVar |--> OutType
84 seCvSubst :: CvSubstEnv, -- InCoVar |--> OutCoercion
85 seIdSubst :: SimplIdSubst, -- InId |--> OutExpr
86
87 ----------- Dynamic part of the environment -----------
88 -- Dynamic in the sense of describing the setup where
89 -- the expression finally ends up
90
91 -- The current set of in-scope variables
92 -- They are all OutVars, and all bound in this module
93 seInScope :: InScopeSet, -- OutVars only
94 -- Includes all variables bound by seFloats
95 seFloats :: Floats,
96 -- See Note [Simplifier floats]
97 seJoinFloats :: JoinFloats
98 -- Handled separately; they don't go very far
99 }
100
101 type StaticEnv = SimplEnv -- Just the static part is relevant
102
103 pprSimplEnv :: SimplEnv -> SDoc
104 -- Used for debugging; selective
105 pprSimplEnv env
106 = vcat [text "TvSubst:" <+> ppr (seTvSubst env),
107 text "CvSubst:" <+> ppr (seCvSubst env),
108 text "IdSubst:" <+> id_subst_doc,
109 text "InScope:" <+> in_scope_vars_doc
110 ]
111 where
112 id_subst_doc = pprUniqFM ppr_id_subst (seIdSubst env)
113 ppr_id_subst (m_ar, sr) = arity_part <+> ppr sr
114 where arity_part = case m_ar of Just ar -> brackets $
115 text "join" <+> int ar
116 Nothing -> empty
117
118 in_scope_vars_doc = pprVarSet (getInScopeVars (seInScope env))
119 (vcat . map ppr_one)
120 ppr_one v | isId v = ppr v <+> ppr (idUnfolding v)
121 | otherwise = ppr v
122
123 type SimplIdSubst = IdEnv (Maybe JoinArity, SimplSR) -- IdId |--> OutExpr
124 -- See Note [Extending the Subst] in CoreSubst
125 -- See Note [Join arity in SimplIdSubst]
126
127 -- | A substitution result.
128 data SimplSR
129 = DoneEx OutExpr -- Completed term
130 | DoneId OutId -- Completed term variable
131 | ContEx TvSubstEnv -- A suspended substitution
132 CvSubstEnv
133 SimplIdSubst
134 InExpr
135
136 instance Outputable SimplSR where
137 ppr (DoneEx e) = text "DoneEx" <+> ppr e
138 ppr (DoneId v) = text "DoneId" <+> ppr v
139 ppr (ContEx _tv _cv _id e) = vcat [text "ContEx" <+> ppr e {-,
140 ppr (filter_env tv), ppr (filter_env id) -}]
141 -- where
142 -- fvs = exprFreeVars e
143 -- filter_env env = filterVarEnv_Directly keep env
144 -- keep uniq _ = uniq `elemUFM_Directly` fvs
145
146 {-
147 Note [SimplEnv invariants]
148 ~~~~~~~~~~~~~~~~~~~~~~~~~~
149 seInScope:
150 The in-scope part of Subst includes *all* in-scope TyVars and Ids
151 The elements of the set may have better IdInfo than the
152 occurrences of in-scope Ids, and (more important) they will
153 have a correctly-substituted type. So we use a lookup in this
154 set to replace occurrences
155
156 The Ids in the InScopeSet are replete with their Rules,
157 and as we gather info about the unfolding of an Id, we replace
158 it in the in-scope set.
159
160 The in-scope set is actually a mapping OutVar -> OutVar, and
161 in case expressions we sometimes bind
162
163 seIdSubst:
164 The substitution is *apply-once* only, because InIds and OutIds
165 can overlap.
166 For example, we generally omit mappings
167 a77 -> a77
168 from the substitution, when we decide not to clone a77, but it's quite
169 legitimate to put the mapping in the substitution anyway.
170
171 Furthermore, consider
172 let x = case k of I# x77 -> ... in
173 let y = case k of I# x77 -> ... in ...
174 and suppose the body is strict in both x and y. Then the simplifier
175 will pull the first (case k) to the top; so the second (case k) will
176 cancel out, mapping x77 to, well, x77! But one is an in-Id and the
177 other is an out-Id.
178
179 Of course, the substitution *must* applied! Things in its domain
180 simply aren't necessarily bound in the result.
181
182 * substId adds a binding (DoneId new_id) to the substitution if
183 the Id's unique has changed
184
185 Note, though that the substitution isn't necessarily extended
186 if the type of the Id changes. Why not? Because of the next point:
187
188 * We *always, always* finish by looking up in the in-scope set
189 any variable that doesn't get a DoneEx or DoneVar hit in the substitution.
190 Reason: so that we never finish up with a "old" Id in the result.
191 An old Id might point to an old unfolding and so on... which gives a space
192 leak.
193
194 [The DoneEx and DoneVar hits map to "new" stuff.]
195
196 * It follows that substExpr must not do a no-op if the substitution is empty.
197 substType is free to do so, however.
198
199 * When we come to a let-binding (say) we generate new IdInfo, including an
200 unfolding, attach it to the binder, and add this newly adorned binder to
201 the in-scope set. So all subsequent occurrences of the binder will get
202 mapped to the full-adorned binder, which is also the one put in the
203 binding site.
204
205 * The in-scope "set" usually maps x->x; we use it simply for its domain.
206 But sometimes we have two in-scope Ids that are synomyms, and should
207 map to the same target: x->x, y->x. Notably:
208 case y of x { ... }
209 That's why the "set" is actually a VarEnv Var
210
211 Note [Join arity in SimplIdSubst]
212 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
213
214 We have to remember which incoming variables are join points (the occurrences
215 may not be marked correctly yet; we're in change of propagating the change if
216 OccurAnal makes something a join point). Normally the in-scope set is where we
217 keep the latest information, but the in-scope set tracks only OutVars; if a
218 binding is unconditionally inlined, it never makes it into the in-scope set,
219 and we need to know at the occurrence site that the variable is a join point so
220 that we know to drop the context. Thus we remember which join points we're
221 substituting. Clumsily, finding whether an InVar is a join variable may require
222 looking in both the substitution *and* the in-scope set (see
223 'isJoinIdInEnv_maybe').
224 -}
225
226 mkSimplEnv :: SimplifierMode -> SimplEnv
227 mkSimplEnv mode
228 = SimplEnv { seMode = mode
229 , seInScope = init_in_scope
230 , seFloats = emptyFloats
231 , seJoinFloats = emptyJoinFloats
232 , seTvSubst = emptyVarEnv
233 , seCvSubst = emptyVarEnv
234 , seIdSubst = emptyVarEnv }
235 -- The top level "enclosing CC" is "SUBSUMED".
236
237 init_in_scope :: InScopeSet
238 init_in_scope = mkInScopeSet (unitVarSet (mkWildValBinder unitTy))
239 -- See Note [WildCard binders]
240
241 {-
242 Note [WildCard binders]
243 ~~~~~~~~~~~~~~~~~~~~~~~
244 The program to be simplified may have wild binders
245 case e of wild { p -> ... }
246 We want to *rename* them away, so that there are no
247 occurrences of 'wild-id' (with wildCardKey). The easy
248 way to do that is to start of with a representative
249 Id in the in-scope set
250
251 There can be be *occurrences* of wild-id. For example,
252 MkCore.mkCoreApp transforms
253 e (a /# b) --> case (a /# b) of wild { DEFAULT -> e wild }
254 This is ok provided 'wild' isn't free in 'e', and that's the delicate
255 thing. Generally, you want to run the simplifier to get rid of the
256 wild-ids before doing much else.
257
258 It's a very dark corner of GHC. Maybe it should be cleaned up.
259 -}
260
261 getMode :: SimplEnv -> SimplifierMode
262 getMode env = seMode env
263
264 setMode :: SimplifierMode -> SimplEnv -> SimplEnv
265 setMode mode env = env { seMode = mode }
266
267 updMode :: (SimplifierMode -> SimplifierMode) -> SimplEnv -> SimplEnv
268 updMode upd env = env { seMode = upd (seMode env) }
269
270 ---------------------
271 extendIdSubst :: SimplEnv -> Id -> SimplSR -> SimplEnv
272 extendIdSubst env@(SimplEnv {seIdSubst = subst}) var res
273 = ASSERT2( isId var && not (isCoVar var), ppr var )
274 env { seIdSubst = extendVarEnv subst var (isJoinId_maybe var, res) }
275
276 extendTvSubst :: SimplEnv -> TyVar -> Type -> SimplEnv
277 extendTvSubst env@(SimplEnv {seTvSubst = tsubst}) var res
278 = ASSERT( isTyVar var )
279 env {seTvSubst = extendVarEnv tsubst var res}
280
281 extendCvSubst :: SimplEnv -> CoVar -> Coercion -> SimplEnv
282 extendCvSubst env@(SimplEnv {seCvSubst = csubst}) var co
283 = ASSERT( isCoVar var )
284 env {seCvSubst = extendVarEnv csubst var co}
285
286 ---------------------
287 getInScope :: SimplEnv -> InScopeSet
288 getInScope env = seInScope env
289
290 setInScopeSet :: SimplEnv -> InScopeSet -> SimplEnv
291 setInScopeSet env in_scope = env {seInScope = in_scope}
292
293 setInScope :: SimplEnv -> SimplEnv -> SimplEnv
294 -- Set the in-scope set, and *zap* the floats
295 setInScope env env_with_scope
296 = env { seInScope = seInScope env_with_scope,
297 seFloats = emptyFloats,
298 seJoinFloats = emptyJoinFloats }
299
300 setFloats :: SimplEnv -> SimplEnv -> SimplEnv
301 -- Set the in-scope set *and* the floats
302 setFloats env env_with_floats
303 = env { seInScope = seInScope env_with_floats,
304 seFloats = seFloats env_with_floats,
305 seJoinFloats = seJoinFloats env_with_floats }
306
307 restoreJoinFloats :: SimplEnv -> SimplEnv -> SimplEnv
308 -- Put back floats previously zapped
309 -- Unlike 'setFloats', does *not* update the in-scope set, since the right-hand
310 -- env is assumed to be *older*
311 restoreJoinFloats env old_env
312 = env { seJoinFloats = seJoinFloats old_env }
313
314 addNewInScopeIds :: SimplEnv -> [CoreBndr] -> SimplEnv
315 -- The new Ids are guaranteed to be freshly allocated
316 addNewInScopeIds env@(SimplEnv { seInScope = in_scope, seIdSubst = id_subst }) vs
317 = env { seInScope = in_scope `extendInScopeSetList` vs,
318 seIdSubst = id_subst `delVarEnvList` vs }
319 -- Why delete? Consider
320 -- let x = a*b in (x, \x -> x+3)
321 -- We add [x |-> a*b] to the substitution, but we must
322 -- _delete_ it from the substitution when going inside
323 -- the (\x -> ...)!
324
325 modifyInScope :: SimplEnv -> CoreBndr -> SimplEnv
326 -- The variable should already be in scope, but
327 -- replace the existing version with this new one
328 -- which has more information
329 modifyInScope env@(SimplEnv {seInScope = in_scope}) v
330 = env {seInScope = extendInScopeSet in_scope v}
331
332 ---------------------
333 zapSubstEnv :: SimplEnv -> SimplEnv
334 zapSubstEnv env = env {seTvSubst = emptyVarEnv, seCvSubst = emptyVarEnv, seIdSubst = emptyVarEnv}
335
336 setSubstEnv :: SimplEnv -> TvSubstEnv -> CvSubstEnv -> SimplIdSubst -> SimplEnv
337 setSubstEnv env tvs cvs ids = env { seTvSubst = tvs, seCvSubst = cvs, seIdSubst = ids }
338
339 mkContEx :: SimplEnv -> InExpr -> SimplSR
340 mkContEx (SimplEnv { seTvSubst = tvs, seCvSubst = cvs, seIdSubst = ids }) e = ContEx tvs cvs ids e
341
342 {-
343 ************************************************************************
344 * *
345 \subsection{Floats}
346 * *
347 ************************************************************************
348
349 Note [Simplifier floats]
350 ~~~~~~~~~~~~~~~~~~~~~~~~~
351 The Floats is a bunch of bindings, classified by a FloatFlag.
352
353 * All of them satisfy the let/app invariant
354
355 Examples
356
357 NonRec x (y:ys) FltLifted
358 Rec [(x,rhs)] FltLifted
359
360 NonRec x* (p:q) FltOKSpec -- RHS is WHNF. Question: why not FltLifted?
361 NonRec x# (y +# 3) FltOkSpec -- Unboxed, but ok-for-spec'n
362
363 NonRec x* (f y) FltCareful -- Strict binding; might fail or diverge
364
365 Can't happen:
366 NonRec x# (a /# b) -- Might fail; does not satisfy let/app
367 NonRec x# (f y) -- Might diverge; does not satisfy let/app
368 -}
369
370 data Floats = Floats (OrdList OutBind) FloatFlag
371 -- See Note [Simplifier floats]
372
373 type JoinFloats = OrdList OutBind
374
375 data FloatFlag
376 = FltLifted -- All bindings are lifted and lazy *or*
377 -- consist of a single primitive string literal
378 -- Hence ok to float to top level, or recursive
379
380 | FltOkSpec -- All bindings are FltLifted *or*
381 -- strict (perhaps because unlifted,
382 -- perhaps because of a strict binder),
383 -- *and* ok-for-speculation
384 -- Hence ok to float out of the RHS
385 -- of a lazy non-recursive let binding
386 -- (but not to top level, or into a rec group)
387
388 | FltCareful -- At least one binding is strict (or unlifted)
389 -- and not guaranteed cheap
390 -- Do not float these bindings out of a lazy let
391
392 instance Outputable Floats where
393 ppr (Floats binds ff) = ppr ff $$ ppr (fromOL binds)
394
395 instance Outputable FloatFlag where
396 ppr FltLifted = text "FltLifted"
397 ppr FltOkSpec = text "FltOkSpec"
398 ppr FltCareful = text "FltCareful"
399
400 andFF :: FloatFlag -> FloatFlag -> FloatFlag
401 andFF FltCareful _ = FltCareful
402 andFF FltOkSpec FltCareful = FltCareful
403 andFF FltOkSpec _ = FltOkSpec
404 andFF FltLifted flt = flt
405
406 doFloatFromRhs :: TopLevelFlag -> RecFlag -> Bool -> OutExpr -> SimplEnv -> Bool
407 -- If you change this function look also at FloatIn.noFloatFromRhs
408 doFloatFromRhs lvl rec str rhs (SimplEnv {seFloats = Floats fs ff})
409 = not (isNilOL fs) && want_to_float && can_float
410 where
411 want_to_float = isTopLevel lvl || exprIsCheap rhs || exprIsExpandable rhs
412 -- See Note [Float when cheap or expandable]
413 can_float = case ff of
414 FltLifted -> True
415 FltOkSpec -> isNotTopLevel lvl && isNonRec rec
416 FltCareful -> isNotTopLevel lvl && isNonRec rec && str
417
418 {-
419 Note [Float when cheap or expandable]
420 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
421 We want to float a let from a let if the residual RHS is
422 a) cheap, such as (\x. blah)
423 b) expandable, such as (f b) if f is CONLIKE
424 But there are
425 - cheap things that are not expandable (eg \x. expensive)
426 - expandable things that are not cheap (eg (f b) where b is CONLIKE)
427 so we must take the 'or' of the two.
428 -}
429
430 emptyFloats :: Floats
431 emptyFloats = Floats nilOL FltLifted
432
433 emptyJoinFloats :: JoinFloats
434 emptyJoinFloats = nilOL
435
436 unitFloat :: OutBind -> Floats
437 -- This key function constructs a singleton float with the right form
438 unitFloat bind = ASSERT(all (not . isJoinId) (bindersOf bind))
439 Floats (unitOL bind) (flag bind)
440 where
441 flag (Rec {}) = FltLifted
442 flag (NonRec bndr rhs)
443 | not (isStrictId bndr) = FltLifted
444 | exprIsLiteralString rhs = FltLifted
445 -- String literals can be floated freely.
446 -- See Note [CoreSyn top-level string ltierals] in CoreSyn.
447 | exprOkForSpeculation rhs = FltOkSpec -- Unlifted, and lifted but ok-for-spec (eg HNF)
448 | otherwise = ASSERT2( not (isUnliftedType (idType bndr)), ppr bndr )
449 FltCareful
450 -- Unlifted binders can only be let-bound if exprOkForSpeculation holds
451
452 unitJoinFloat :: OutBind -> JoinFloats
453 unitJoinFloat bind = ASSERT(all isJoinId (bindersOf bind))
454 unitOL bind
455
456 addNonRec :: SimplEnv -> OutId -> OutExpr -> SimplEnv
457 -- Add a non-recursive binding and extend the in-scope set
458 -- The latter is important; the binder may already be in the
459 -- in-scope set (although it might also have been created with newId)
460 -- but it may now have more IdInfo
461 addNonRec env id rhs
462 = id `seq` -- This seq forces the Id, and hence its IdInfo,
463 -- and hence any inner substitutions
464 env { seFloats = floats',
465 seJoinFloats = jfloats',
466 seInScope = extendInScopeSet (seInScope env) id }
467 where
468 bind = NonRec id rhs
469
470 floats' | isJoinId id = seFloats env
471 | otherwise = seFloats env `addFlts` unitFloat bind
472 jfloats' | isJoinId id = seJoinFloats env `addJoinFlts` unitJoinFloat bind
473 | otherwise = seJoinFloats env
474
475 extendFloats :: SimplEnv -> OutBind -> SimplEnv
476 -- Add these bindings to the floats, and extend the in-scope env too
477 extendFloats env bind
478 = ASSERT(all (not . isJoinId) (bindersOf bind))
479 env { seFloats = floats',
480 seJoinFloats = jfloats',
481 seInScope = extendInScopeSetList (seInScope env) bndrs }
482 where
483 bndrs = bindersOf bind
484
485 floats' | isJoinBind bind = seFloats env
486 | otherwise = seFloats env `addFlts` unitFloat bind
487 jfloats' | isJoinBind bind = seJoinFloats env `addJoinFlts`
488 unitJoinFloat bind
489 | otherwise = seJoinFloats env
490
491 addFloats :: SimplEnv -> SimplEnv -> SimplEnv
492 -- Add the floats for env2 to env1;
493 -- *plus* the in-scope set for env2, which is bigger
494 -- than that for env1
495 addFloats env1 env2
496 = env1 {seFloats = seFloats env1 `addFlts` seFloats env2,
497 seJoinFloats = seJoinFloats env1 `addJoinFlts` seJoinFloats env2,
498 seInScope = seInScope env2 }
499
500 addFlts :: Floats -> Floats -> Floats
501 addFlts (Floats bs1 l1) (Floats bs2 l2)
502 = Floats (bs1 `appOL` bs2) (l1 `andFF` l2)
503
504 addJoinFlts :: JoinFloats -> JoinFloats -> JoinFloats
505 addJoinFlts = appOL
506
507 zapFloats :: SimplEnv -> SimplEnv
508 zapFloats env = env { seFloats = emptyFloats
509 , seJoinFloats = emptyJoinFloats }
510
511 zapJoinFloats :: SimplEnv -> SimplEnv
512 zapJoinFloats env = env { seJoinFloats = emptyJoinFloats }
513
514 addRecFloats :: SimplEnv -> SimplEnv -> SimplEnv
515 -- Flattens the floats from env2 into a single Rec group,
516 -- prepends the floats from env1, and puts the result back in env2
517 -- This is all very specific to the way recursive bindings are
518 -- handled; see Simplify.simplRecBind
519 addRecFloats env1 env2@(SimplEnv {seFloats = Floats bs ff
520 ,seJoinFloats = jbs })
521 = ASSERT2( case ff of { FltLifted -> True; _ -> False }, ppr (fromOL bs) )
522 env2 {seFloats = seFloats env1 `addFlts` floats'
523 ,seJoinFloats = seJoinFloats env1 `addJoinFlts` jfloats'}
524 where
525 floats' | isNilOL bs = emptyFloats
526 | otherwise = unitFloat (Rec (flattenBinds (fromOL bs)))
527 jfloats' | isNilOL jbs = emptyJoinFloats
528 | otherwise = unitJoinFloat (Rec (flattenBinds (fromOL jbs)))
529
530 wrapFloats :: SimplEnv -> OutExpr -> OutExpr
531 -- Wrap the floats around the expression; they should all
532 -- satisfy the let/app invariant, so mkLets should do the job just fine
533 wrapFloats env@(SimplEnv {seFloats = Floats bs _}) body
534 = foldrOL Let (wrapJoinFloats env body) bs
535 -- Note: Always safe to put the joins on the inside since the values
536 -- can't refer to them
537
538 wrapJoinFloats :: SimplEnv -> OutExpr -> OutExpr
539 wrapJoinFloats (SimplEnv {seJoinFloats = jbs}) body
540 = foldrOL Let body jbs
541
542 getFloatBinds :: SimplEnv -> [CoreBind]
543 getFloatBinds env@(SimplEnv {seFloats = Floats bs _})
544 = fromOL bs ++ getJoinFloatBinds env
545
546 getJoinFloatBinds :: SimplEnv -> [CoreBind]
547 getJoinFloatBinds (SimplEnv {seJoinFloats = jbs})
548 = fromOL jbs
549
550 isEmptyFloats :: SimplEnv -> Bool
551 isEmptyFloats env@(SimplEnv {seFloats = Floats bs _})
552 = isNilOL bs && isEmptyJoinFloats env
553
554 isEmptyJoinFloats :: SimplEnv -> Bool
555 isEmptyJoinFloats (SimplEnv {seJoinFloats = jbs})
556 = isNilOL jbs
557
558 mapFloats :: SimplEnv -> ((Id,CoreExpr) -> (Id,CoreExpr)) -> SimplEnv
559 mapFloats env@SimplEnv { seFloats = Floats fs ff, seJoinFloats = jfs } fun
560 = env { seFloats = Floats (mapOL app fs) ff
561 , seJoinFloats = mapOL app jfs }
562 where
563 app (NonRec b e) = case fun (b,e) of (b',e') -> NonRec b' e'
564 app (Rec bs) = Rec (map fun bs)
565
566 {-
567 ************************************************************************
568 * *
569 Substitution of Vars
570 * *
571 ************************************************************************
572
573 Note [Global Ids in the substitution]
574 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
575 We look up even a global (eg imported) Id in the substitution. Consider
576 case X.g_34 of b { (a,b) -> ... case X.g_34 of { (p,q) -> ...} ... }
577 The binder-swap in the occurrence analyser will add a binding
578 for a LocalId version of g (with the same unique though):
579 case X.g_34 of b { (a,b) -> let g_34 = b in
580 ... case X.g_34 of { (p,q) -> ...} ... }
581 So we want to look up the inner X.g_34 in the substitution, where we'll
582 find that it has been substituted by b. (Or conceivably cloned.)
583 -}
584
585 substId :: SimplEnv -> InId -> SimplSR
586 -- Returns DoneEx only on a non-Var expression
587 substId (SimplEnv { seInScope = in_scope, seIdSubst = ids }) v
588 = case snd <$> lookupVarEnv ids v of -- Note [Global Ids in the substitution]
589 Nothing -> DoneId (refineFromInScope in_scope v)
590 Just (DoneId v) -> DoneId (refineFromInScope in_scope v)
591 Just (DoneEx (Var v)) -> DoneId (refineFromInScope in_scope v)
592 Just res -> res -- DoneEx non-var, or ContEx
593
594 -- Get the most up-to-date thing from the in-scope set
595 -- Even though it isn't in the substitution, it may be in
596 -- the in-scope set with better IdInfo
597
598 isJoinIdInEnv_maybe :: SimplEnv -> InId -> Maybe JoinArity
599 isJoinIdInEnv_maybe (SimplEnv { seInScope = inScope, seIdSubst = ids }) v
600 | not (isLocalId v) = Nothing
601 | Just (m_ar, _) <- lookupVarEnv ids v = m_ar
602 | Just v' <- lookupInScope inScope v = isJoinId_maybe v'
603 | otherwise = WARN( True , ppr v )
604 isJoinId_maybe v
605
606 refineFromInScope :: InScopeSet -> Var -> Var
607 refineFromInScope in_scope v
608 | isLocalId v = case lookupInScope in_scope v of
609 Just v' -> v'
610 Nothing -> WARN( True, ppr v ) v -- This is an error!
611 | otherwise = v
612
613 lookupRecBndr :: SimplEnv -> InId -> OutId
614 -- Look up an Id which has been put into the envt by simplRecBndrs,
615 -- but where we have not yet done its RHS
616 lookupRecBndr (SimplEnv { seInScope = in_scope, seIdSubst = ids }) v
617 = case lookupVarEnv ids v of
618 Just (_, DoneId v) -> v
619 Just _ -> pprPanic "lookupRecBndr" (ppr v)
620 Nothing -> refineFromInScope in_scope v
621
622 {-
623 ************************************************************************
624 * *
625 \section{Substituting an Id binder}
626 * *
627 ************************************************************************
628
629
630 These functions are in the monad only so that they can be made strict via seq.
631 -}
632
633 simplBinders :: SimplEnv -> [InBndr] -> SimplM (SimplEnv, [OutBndr])
634 simplBinders env bndrs = mapAccumLM simplBinder env bndrs
635
636 -------------
637 simplBinder :: SimplEnv -> InBndr -> SimplM (SimplEnv, OutBndr)
638 -- Used for lambda and case-bound variables
639 -- Clone Id if necessary, substitute type
640 -- Return with IdInfo already substituted, but (fragile) occurrence info zapped
641 -- The substitution is extended only if the variable is cloned, because
642 -- we *don't* need to use it to track occurrence info.
643 simplBinder env bndr
644 | isTyVar bndr = do { let (env', tv) = substTyVarBndr env bndr
645 ; seqTyVar tv `seq` return (env', tv) }
646 | otherwise = do { let (env', id) = substIdBndr Nothing env bndr
647 ; seqId id `seq` return (env', id) }
648
649 ---------------
650 simplNonRecBndr :: SimplEnv -> InBndr -> SimplM (SimplEnv, OutBndr)
651 -- A non-recursive let binder
652 simplNonRecBndr env id
653 = do { let (env1, id1) = substIdBndr Nothing env id
654 ; seqId id1 `seq` return (env1, id1) }
655
656 ---------------
657 simplNonRecJoinBndr :: SimplEnv -> OutType -> InBndr
658 -> SimplM (SimplEnv, OutBndr)
659 -- A non-recursive let binder for a join point; context being pushed inward may
660 -- change the type
661 simplNonRecJoinBndr env res_ty id
662 = do { let (env1, id1) = substIdBndr (Just res_ty) env id
663 ; seqId id1 `seq` return (env1, id1) }
664
665 ---------------
666 simplRecBndrs :: SimplEnv -> [InBndr] -> SimplM SimplEnv
667 -- Recursive let binders
668 simplRecBndrs env@(SimplEnv {}) ids
669 = ASSERT(all (not . isJoinId) ids)
670 do { let (env1, ids1) = mapAccumL (substIdBndr Nothing) env ids
671 ; seqIds ids1 `seq` return env1 }
672
673 ---------------
674 simplRecJoinBndrs :: SimplEnv -> OutType -> [InBndr] -> SimplM SimplEnv
675 -- Recursive let binders for join points; context being pushed inward may
676 -- change types
677 simplRecJoinBndrs env@(SimplEnv {}) res_ty ids
678 = ASSERT(all isJoinId ids)
679 do { let (env1, ids1) = mapAccumL (substIdBndr (Just res_ty)) env ids
680 ; seqIds ids1 `seq` return env1 }
681
682 ---------------
683 substIdBndr :: Maybe OutType -> SimplEnv -> InBndr -> (SimplEnv, OutBndr)
684 -- Might be a coercion variable
685 substIdBndr new_res_ty env bndr
686 | isCoVar bndr = substCoVarBndr env bndr
687 | otherwise = substNonCoVarIdBndr new_res_ty env bndr
688
689 ---------------
690 substNonCoVarIdBndr
691 :: Maybe OutType -- New result type, if a join binder
692 -> SimplEnv
693 -> InBndr -- Env and binder to transform
694 -> (SimplEnv, OutBndr)
695 -- Clone Id if necessary, substitute its type
696 -- Return an Id with its
697 -- * Type substituted
698 -- * UnfoldingInfo, Rules, WorkerInfo zapped
699 -- * Fragile OccInfo (only) zapped: Note [Robust OccInfo]
700 -- * Robust info, retained especially arity and demand info,
701 -- so that they are available to occurrences that occur in an
702 -- earlier binding of a letrec
703 --
704 -- For the robust info, see Note [Arity robustness]
705 --
706 -- Augment the substitution if the unique changed
707 -- Extend the in-scope set with the new Id
708 --
709 -- Similar to CoreSubst.substIdBndr, except that
710 -- the type of id_subst differs
711 -- all fragile info is zapped
712 substNonCoVarIdBndr new_res_ty
713 env@(SimplEnv { seInScope = in_scope
714 , seIdSubst = id_subst })
715 old_id
716 = ASSERT2( not (isCoVar old_id), ppr old_id )
717 (env { seInScope = in_scope `extendInScopeSet` new_id,
718 seIdSubst = new_subst }, new_id)
719 where
720 id1 = uniqAway in_scope old_id
721 id2 = substIdType env id1
722 id3 | Just res_ty <- new_res_ty
723 = id2 `setIdType` setJoinResTy (idJoinArity id2) res_ty (idType id2)
724 | otherwise
725 = id2
726 new_id = zapFragileIdInfo id3 -- Zaps rules, worker-info, unfolding
727 -- and fragile OccInfo
728
729 -- Extend the substitution if the unique has changed,
730 -- or there's some useful occurrence information
731 -- See the notes with substTyVarBndr for the delSubstEnv
732 new_subst | new_id /= old_id
733 = extendVarEnv id_subst old_id
734 (isJoinId_maybe new_id, DoneId new_id)
735 | otherwise
736 = delVarEnv id_subst old_id
737
738 ------------------------------------
739 seqTyVar :: TyVar -> ()
740 seqTyVar b = b `seq` ()
741
742 seqId :: Id -> ()
743 seqId id = seqType (idType id) `seq`
744 idInfo id `seq`
745 ()
746
747 seqIds :: [Id] -> ()
748 seqIds [] = ()
749 seqIds (id:ids) = seqId id `seq` seqIds ids
750
751 {-
752 Note [Arity robustness]
753 ~~~~~~~~~~~~~~~~~~~~~~~
754 We *do* transfer the arity from from the in_id of a let binding to the
755 out_id. This is important, so that the arity of an Id is visible in
756 its own RHS. For example:
757 f = \x. ....g (\y. f y)....
758 We can eta-reduce the arg to g, because f is a value. But that
759 needs to be visible.
760
761 This interacts with the 'state hack' too:
762 f :: Bool -> IO Int
763 f = \x. case x of
764 True -> f y
765 False -> \s -> ...
766 Can we eta-expand f? Only if we see that f has arity 1, and then we
767 take advantage of the 'state hack' on the result of
768 (f y) :: State# -> (State#, Int) to expand the arity one more.
769
770 There is a disadvantage though. Making the arity visible in the RHS
771 allows us to eta-reduce
772 f = \x -> f x
773 to
774 f = f
775 which technically is not sound. This is very much a corner case, so
776 I'm not worried about it. Another idea is to ensure that f's arity
777 never decreases; its arity started as 1, and we should never eta-reduce
778 below that.
779
780
781 Note [Robust OccInfo]
782 ~~~~~~~~~~~~~~~~~~~~~
783 It's important that we *do* retain the loop-breaker OccInfo, because
784 that's what stops the Id getting inlined infinitely, in the body of
785 the letrec.
786 -}
787
788
789 {-
790 ************************************************************************
791 * *
792 Impedence matching to type substitution
793 * *
794 ************************************************************************
795 -}
796
797 getTCvSubst :: SimplEnv -> TCvSubst
798 getTCvSubst (SimplEnv { seInScope = in_scope, seTvSubst = tv_env
799 , seCvSubst = cv_env })
800 = mkTCvSubst in_scope (tv_env, cv_env)
801
802 substTy :: SimplEnv -> Type -> Type
803 substTy env ty = Type.substTy (getTCvSubst env) ty
804
805 substTyVar :: SimplEnv -> TyVar -> Type
806 substTyVar env tv = Type.substTyVar (getTCvSubst env) tv
807
808 substTyVarBndr :: SimplEnv -> TyVar -> (SimplEnv, TyVar)
809 substTyVarBndr env tv
810 = case Type.substTyVarBndr (getTCvSubst env) tv of
811 (TCvSubst in_scope' tv_env' cv_env', tv')
812 -> (env { seInScope = in_scope', seTvSubst = tv_env', seCvSubst = cv_env' }, tv')
813
814 substCoVar :: SimplEnv -> CoVar -> Coercion
815 substCoVar env tv = Coercion.substCoVar (getTCvSubst env) tv
816
817 substCoVarBndr :: SimplEnv -> CoVar -> (SimplEnv, CoVar)
818 substCoVarBndr env cv
819 = case Coercion.substCoVarBndr (getTCvSubst env) cv of
820 (TCvSubst in_scope' tv_env' cv_env', cv')
821 -> (env { seInScope = in_scope', seTvSubst = tv_env', seCvSubst = cv_env' }, cv')
822
823 substCo :: SimplEnv -> Coercion -> Coercion
824 substCo env co = Coercion.substCo (getTCvSubst env) co
825
826 ------------------
827 substIdType :: SimplEnv -> Id -> Id
828 substIdType (SimplEnv { seInScope = in_scope, seTvSubst = tv_env, seCvSubst = cv_env }) id
829 | (isEmptyVarEnv tv_env && isEmptyVarEnv cv_env)
830 || noFreeVarsOfType old_ty
831 = id
832 | otherwise = Id.setIdType id (Type.substTy (TCvSubst in_scope tv_env cv_env) old_ty)
833 -- The tyCoVarsOfType is cheaper than it looks
834 -- because we cache the free tyvars of the type
835 -- in a Note in the id's type itself
836 where
837 old_ty = idType id