4240647d581262c63117222b12354aa1cfab3bca
[ghc.git] / compiler / coreSyn / CoreOpt.hs
1 {-
2 (c) The University of Glasgow 2006
3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 -}
5
6 {-# LANGUAGE CPP #-}
7 module CoreOpt (
8 -- ** Simple expression optimiser
9 simpleOptPgm, simpleOptExpr, simpleOptExprWith,
10
11 -- ** Join points
12 joinPointBinding_maybe, joinPointBindings_maybe,
13
14 -- ** Predicates on expressions
15 exprIsConApp_maybe, exprIsLiteral_maybe, exprIsLambda_maybe,
16
17 -- ** Coercions and casts
18 pushCoArg, pushCoValArg, pushCoTyArg, collectBindersPushingCo
19 ) where
20
21 #include "HsVersions.h"
22
23 import GhcPrelude
24
25 import CoreArity( joinRhsArity, etaExpandToJoinPoint )
26
27 import CoreSyn
28 import CoreSubst
29 import CoreUtils
30 import CoreFVs
31 import PprCore ( pprCoreBindings, pprRules )
32 import OccurAnal( occurAnalyseExpr, occurAnalysePgm )
33 import Literal ( Literal(MachStr) )
34 import Id
35 import Var ( varType )
36 import VarSet
37 import VarEnv
38 import DataCon
39 import OptCoercion ( optCoercion )
40 import Type hiding ( substTy, extendTvSubst, extendCvSubst, extendTvSubstList
41 , isInScope, substTyVarBndr, cloneTyVarBndr )
42 import Coercion hiding ( substCo, substCoVarBndr )
43 import TyCon ( tyConArity )
44 import TysWiredIn
45 import PrelNames
46 import BasicTypes
47 import Module ( Module )
48 import ErrUtils
49 import DynFlags
50 import Outputable
51 import Pair
52 import Util
53 import Maybes ( orElse )
54 import FastString
55 import Data.List
56 import qualified Data.ByteString as BS
57
58 {-
59 ************************************************************************
60 * *
61 The Simple Optimiser
62 * *
63 ************************************************************************
64
65 Note [The simple optimiser]
66 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
67 The simple optimiser is a lightweight, pure (non-monadic) function
68 that rapidly does a lot of simple optimisations, including
69
70 - inlining things that occur just once,
71 or whose RHS turns out to be trivial
72 - beta reduction
73 - case of known constructor
74 - dead code elimination
75
76 It does NOT do any call-site inlining; it only inlines a function if
77 it can do so unconditionally, dropping the binding. It thereby
78 guarantees to leave no un-reduced beta-redexes.
79
80 It is careful to follow the guidance of "Secrets of the GHC inliner",
81 and in particular the pre-inline-unconditionally and
82 post-inline-unconditionally story, to do effective beta reduction on
83 functions called precisely once, without repeatedly optimising the same
84 expression. In fact, the simple optimiser is a good example of this
85 little dance in action; the full Simplifier is a lot more complicated.
86
87 -}
88
89 simpleOptExpr :: CoreExpr -> CoreExpr
90 -- See Note [The simple optimiser]
91 -- Do simple optimisation on an expression
92 -- The optimisation is very straightforward: just
93 -- inline non-recursive bindings that are used only once,
94 -- or where the RHS is trivial
95 --
96 -- We also inline bindings that bind a Eq# box: see
97 -- See Note [Getting the map/coerce RULE to work].
98 --
99 -- Also we convert functions to join points where possible (as
100 -- the occurrence analyser does most of the work anyway).
101 --
102 -- The result is NOT guaranteed occurrence-analysed, because
103 -- in (let x = y in ....) we substitute for x; so y's occ-info
104 -- may change radically
105
106 simpleOptExpr expr
107 = -- pprTrace "simpleOptExpr" (ppr init_subst $$ ppr expr)
108 simpleOptExprWith init_subst expr
109 where
110 init_subst = mkEmptySubst (mkInScopeSet (exprFreeVars expr))
111 -- It's potentially important to make a proper in-scope set
112 -- Consider let x = ..y.. in \y. ...x...
113 -- Then we should remember to clone y before substituting
114 -- for x. It's very unlikely to occur, because we probably
115 -- won't *be* substituting for x if it occurs inside a
116 -- lambda.
117 --
118 -- It's a bit painful to call exprFreeVars, because it makes
119 -- three passes instead of two (occ-anal, and go)
120
121 simpleOptExprWith :: Subst -> InExpr -> OutExpr
122 -- See Note [The simple optimiser]
123 simpleOptExprWith subst expr
124 = simple_opt_expr init_env (occurAnalyseExpr expr)
125 where
126 init_env = SOE { soe_inl = emptyVarEnv, soe_subst = subst }
127
128 ----------------------
129 simpleOptPgm :: DynFlags -> Module
130 -> CoreProgram -> [CoreRule] -> [CoreVect]
131 -> IO (CoreProgram, [CoreRule], [CoreVect])
132 -- See Note [The simple optimiser]
133 simpleOptPgm dflags this_mod binds rules vects
134 = do { dumpIfSet_dyn dflags Opt_D_dump_occur_anal "Occurrence analysis"
135 (pprCoreBindings occ_anald_binds $$ pprRules rules );
136
137 ; return (reverse binds', rules', vects') }
138 where
139 occ_anald_binds = occurAnalysePgm this_mod
140 (\_ -> True) {- All unfoldings active -}
141 (\_ -> False) {- No rules active -}
142 rules vects emptyVarSet binds
143
144 (final_env, binds') = foldl do_one (emptyEnv, []) occ_anald_binds
145 final_subst = soe_subst final_env
146
147 rules' = substRulesForImportedIds final_subst rules
148 vects' = substVects final_subst vects
149 -- We never unconditionally inline into rules,
150 -- hence paying just a substitution
151
152 do_one (env, binds') bind
153 = case simple_opt_bind env bind of
154 (env', Nothing) -> (env', binds')
155 (env', Just bind') -> (env', bind':binds')
156
157 -- In these functions the substitution maps InVar -> OutExpr
158
159 ----------------------
160 type SimpleClo = (SimpleOptEnv, InExpr)
161
162 data SimpleOptEnv
163 = SOE { soe_inl :: IdEnv SimpleClo
164 -- Deals with preInlineUnconditionally; things
165 -- that occur exactly once and are inlined
166 -- without having first been simplified
167
168 , soe_subst :: Subst
169 -- Deals with cloning; includes the InScopeSet
170 }
171
172 instance Outputable SimpleOptEnv where
173 ppr (SOE { soe_inl = inl, soe_subst = subst })
174 = text "SOE {" <+> vcat [ text "soe_inl =" <+> ppr inl
175 , text "soe_subst =" <+> ppr subst ]
176 <+> text "}"
177
178 emptyEnv :: SimpleOptEnv
179 emptyEnv = SOE { soe_inl = emptyVarEnv
180 , soe_subst = emptySubst }
181
182 soeZapSubst :: SimpleOptEnv -> SimpleOptEnv
183 soeZapSubst (SOE { soe_subst = subst })
184 = SOE { soe_inl = emptyVarEnv, soe_subst = zapSubstEnv subst }
185
186 soeSetInScope :: SimpleOptEnv -> SimpleOptEnv -> SimpleOptEnv
187 -- Take in-scope set from env1, and the rest from env2
188 soeSetInScope (SOE { soe_subst = subst1 })
189 env2@(SOE { soe_subst = subst2 })
190 = env2 { soe_subst = setInScope subst2 (substInScope subst1) }
191
192 ---------------
193 simple_opt_clo :: SimpleOptEnv -> SimpleClo -> OutExpr
194 simple_opt_clo env (e_env, e)
195 = simple_opt_expr (soeSetInScope env e_env) e
196
197 simple_opt_expr :: SimpleOptEnv -> InExpr -> OutExpr
198 simple_opt_expr env expr
199 = go expr
200 where
201 subst = soe_subst env
202 in_scope = substInScope subst
203 in_scope_env = (in_scope, simpleUnfoldingFun)
204
205 go (Var v)
206 | Just clo <- lookupVarEnv (soe_inl env) v
207 = simple_opt_clo env clo
208 | otherwise
209 = lookupIdSubst (text "simpleOptExpr") (soe_subst env) v
210
211 go (App e1 e2) = simple_app env e1 [(env,e2)]
212 go (Type ty) = Type (substTy subst ty)
213 go (Coercion co) = Coercion (optCoercion (getTCvSubst subst) co)
214 go (Lit lit) = Lit lit
215 go (Tick tickish e) = mkTick (substTickish subst tickish) (go e)
216 go (Cast e co) | isReflCo co' = go e
217 | otherwise = Cast (go e) co'
218 where
219 co' = optCoercion (getTCvSubst subst) co
220
221 go (Let bind body) = case simple_opt_bind env bind of
222 (env', Nothing) -> simple_opt_expr env' body
223 (env', Just bind) -> Let bind (simple_opt_expr env' body)
224
225 go lam@(Lam {}) = go_lam env [] lam
226 go (Case e b ty as)
227 -- See Note [Getting the map/coerce RULE to work]
228 | isDeadBinder b
229 , Just (con, _tys, es) <- exprIsConApp_maybe in_scope_env e'
230 , Just (altcon, bs, rhs) <- findAlt (DataAlt con) as
231 = case altcon of
232 DEFAULT -> go rhs
233 _ -> foldr wrapLet (simple_opt_expr env' rhs) mb_prs
234 where
235 (env', mb_prs) = mapAccumL simple_out_bind env $
236 zipEqual "simpleOptExpr" bs es
237
238 -- Note [Getting the map/coerce RULE to work]
239 | isDeadBinder b
240 , [(DEFAULT, _, rhs)] <- as
241 , isCoercionType (varType b)
242 , (Var fun, _args) <- collectArgs e
243 , fun `hasKey` coercibleSCSelIdKey
244 -- without this last check, we get #11230
245 = go rhs
246
247 | otherwise
248 = Case e' b' (substTy subst ty)
249 (map (go_alt env') as)
250 where
251 e' = go e
252 (env', b') = subst_opt_bndr env b
253
254 ----------------------
255 go_alt env (con, bndrs, rhs)
256 = (con, bndrs', simple_opt_expr env' rhs)
257 where
258 (env', bndrs') = subst_opt_bndrs env bndrs
259
260 ----------------------
261 -- go_lam tries eta reduction
262 go_lam env bs' (Lam b e)
263 = go_lam env' (b':bs') e
264 where
265 (env', b') = subst_opt_bndr env b
266 go_lam env bs' e
267 | Just etad_e <- tryEtaReduce bs e' = etad_e
268 | otherwise = mkLams bs e'
269 where
270 bs = reverse bs'
271 e' = simple_opt_expr env e
272
273 ----------------------
274 -- simple_app collects arguments for beta reduction
275 simple_app :: SimpleOptEnv -> InExpr -> [SimpleClo] -> CoreExpr
276
277 simple_app env (Var v) as
278 | Just (env', e) <- lookupVarEnv (soe_inl env) v
279 = simple_app (soeSetInScope env env') e as
280
281 | let unf = idUnfolding v
282 , isCompulsoryUnfolding (idUnfolding v)
283 , isAlwaysActive (idInlineActivation v)
284 -- See Note [Unfold compulsory unfoldings in LHSs]
285 = simple_app (soeZapSubst env) (unfoldingTemplate unf) as
286
287 | otherwise
288 , let out_fn = lookupIdSubst (text "simple_app") (soe_subst env) v
289 = finish_app env out_fn as
290
291 simple_app env (App e1 e2) as
292 = simple_app env e1 ((env, e2) : as)
293
294 simple_app env (Lam b e) (a:as)
295 = wrapLet mb_pr (simple_app env' e as)
296 where
297 (env', mb_pr) = simple_bind_pair env b Nothing a
298
299 simple_app env (Tick t e) as
300 -- Okay to do "(Tick t e) x ==> Tick t (e x)"?
301 | t `tickishScopesLike` SoftScope
302 = mkTick t $ simple_app env e as
303
304 simple_app env e as
305 = finish_app env (simple_opt_expr env e) as
306
307 finish_app :: SimpleOptEnv -> OutExpr -> [SimpleClo] -> OutExpr
308 finish_app _ fun []
309 = fun
310 finish_app env fun (arg:args)
311 = finish_app env (App fun (simple_opt_clo env arg)) args
312
313 ----------------------
314 simple_opt_bind :: SimpleOptEnv -> InBind
315 -> (SimpleOptEnv, Maybe OutBind)
316 simple_opt_bind env (NonRec b r)
317 = (env', case mb_pr of
318 Nothing -> Nothing
319 Just (b,r) -> Just (NonRec b r))
320 where
321 (b', r') = joinPointBinding_maybe b r `orElse` (b, r)
322 (env', mb_pr) = simple_bind_pair env b' Nothing (env,r')
323
324 simple_opt_bind env (Rec prs)
325 = (env'', res_bind)
326 where
327 res_bind = Just (Rec (reverse rev_prs'))
328 prs' = joinPointBindings_maybe prs `orElse` prs
329 (env', bndrs') = subst_opt_bndrs env (map fst prs')
330 (env'', rev_prs') = foldl do_pr (env', []) (prs' `zip` bndrs')
331 do_pr (env, prs) ((b,r), b')
332 = (env', case mb_pr of
333 Just pr -> pr : prs
334 Nothing -> prs)
335 where
336 (env', mb_pr) = simple_bind_pair env b (Just b') (env,r)
337
338 ----------------------
339 simple_bind_pair :: SimpleOptEnv
340 -> InVar -> Maybe OutVar
341 -> SimpleClo
342 -> (SimpleOptEnv, Maybe (OutVar, OutExpr))
343 -- (simple_bind_pair subst in_var out_rhs)
344 -- either extends subst with (in_var -> out_rhs)
345 -- or returns Nothing
346 simple_bind_pair env@(SOE { soe_inl = inl_env, soe_subst = subst })
347 in_bndr mb_out_bndr clo@(rhs_env, in_rhs)
348 | Type ty <- in_rhs -- let a::* = TYPE ty in <body>
349 , let out_ty = substTy (soe_subst rhs_env) ty
350 = ASSERT( isTyVar in_bndr )
351 (env { soe_subst = extendTvSubst subst in_bndr out_ty }, Nothing)
352
353 | Coercion co <- in_rhs
354 , let out_co = optCoercion (getTCvSubst (soe_subst rhs_env)) co
355 = ASSERT( isCoVar in_bndr )
356 (env { soe_subst = extendCvSubst subst in_bndr out_co }, Nothing)
357
358 | pre_inline_unconditionally
359 = (env { soe_inl = extendVarEnv inl_env in_bndr clo }, Nothing)
360
361 | otherwise
362 = simple_out_bind_pair env in_bndr mb_out_bndr
363 (simple_opt_clo env clo)
364 occ active stable_unf
365 where
366 stable_unf = isStableUnfolding (idUnfolding in_bndr)
367 active = isAlwaysActive (idInlineActivation in_bndr)
368 occ = idOccInfo in_bndr
369
370 pre_inline_unconditionally :: Bool
371 pre_inline_unconditionally
372 | isCoVar in_bndr = False -- See Note [Do not inline CoVars unconditionally]
373 | isExportedId in_bndr = False -- in SimplUtils
374 | stable_unf = False
375 | not active = False -- Note [Inline prag in simplOpt]
376 | not (safe_to_inline occ) = False
377 | otherwise = True
378
379 -- Unconditionally safe to inline
380 safe_to_inline :: OccInfo -> Bool
381 safe_to_inline (IAmALoopBreaker {}) = False
382 safe_to_inline IAmDead = True
383 safe_to_inline occ@(OneOcc {}) = not (occ_in_lam occ)
384 && occ_one_br occ
385 safe_to_inline (ManyOccs {}) = False
386
387 -------------------
388 simple_out_bind :: SimpleOptEnv -> (InVar, OutExpr)
389 -> (SimpleOptEnv, Maybe (OutVar, OutExpr))
390 simple_out_bind env@(SOE { soe_subst = subst }) (in_bndr, out_rhs)
391 | Type out_ty <- out_rhs
392 = ASSERT( isTyVar in_bndr )
393 (env { soe_subst = extendTvSubst subst in_bndr out_ty }, Nothing)
394
395 | Coercion out_co <- out_rhs
396 = ASSERT( isCoVar in_bndr )
397 (env { soe_subst = extendCvSubst subst in_bndr out_co }, Nothing)
398
399 | otherwise
400 = simple_out_bind_pair env in_bndr Nothing out_rhs
401 (idOccInfo in_bndr) True False
402
403 -------------------
404 simple_out_bind_pair :: SimpleOptEnv
405 -> InId -> Maybe OutId -> OutExpr
406 -> OccInfo -> Bool -> Bool
407 -> (SimpleOptEnv, Maybe (OutVar, OutExpr))
408 simple_out_bind_pair env in_bndr mb_out_bndr out_rhs
409 occ_info active stable_unf
410 | post_inline_unconditionally
411 = ( env' { soe_subst = extendIdSubst (soe_subst env) in_bndr out_rhs }
412 , Nothing)
413
414 | otherwise
415 = ( env', Just (out_bndr, out_rhs) )
416 where
417 (env', bndr1) = case mb_out_bndr of
418 Just out_bndr -> (env, out_bndr)
419 Nothing -> subst_opt_bndr env in_bndr
420 out_bndr = add_info env' in_bndr bndr1
421
422 post_inline_unconditionally :: Bool
423 post_inline_unconditionally
424 | not active = False
425 | isWeakLoopBreaker occ_info = False -- If it's a loop-breaker of any kind, don't inline
426 -- because it might be referred to "earlier"
427 | stable_unf = False -- Note [Stable unfoldings and postInlineUnconditionally]
428 | isExportedId in_bndr = False -- Note [Exported Ids and trivial RHSs]
429 | exprIsTrivial out_rhs = True
430 | coercible_hack = True
431 | otherwise = False
432
433 -- See Note [Getting the map/coerce RULE to work]
434 coercible_hack | (Var fun, args) <- collectArgs out_rhs
435 , Just dc <- isDataConWorkId_maybe fun
436 , dc `hasKey` heqDataConKey || dc `hasKey` coercibleDataConKey
437 = all exprIsTrivial args
438 | otherwise
439 = False
440
441 {- Note [Exported Ids and trivial RHSs]
442 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
443 We obviously do not want to unconditionally inline an Id that is exported.
444 In SimplUtils, Note [Top level and postInlineUnconditionally], we
445 explain why we don't inline /any/ top-level things unconditionally, even
446 trivial ones. But we do here! Why? In the simple optimiser
447
448 * We do no rule rewrites
449 * We do no call-site inlining
450
451 Those differences obviate the reasons for not inlining a trivial rhs,
452 and increase the benefit for doing so. So we unconditionally inline trivial
453 rhss here.
454 -}
455
456 ----------------------
457 subst_opt_bndrs :: SimpleOptEnv -> [InVar] -> (SimpleOptEnv, [OutVar])
458 subst_opt_bndrs env bndrs = mapAccumL subst_opt_bndr env bndrs
459
460 subst_opt_bndr :: SimpleOptEnv -> InVar -> (SimpleOptEnv, OutVar)
461 subst_opt_bndr env bndr
462 | isTyVar bndr = (env { soe_subst = subst_tv }, tv')
463 | isCoVar bndr = (env { soe_subst = subst_cv }, cv')
464 | otherwise = subst_opt_id_bndr env bndr
465 where
466 subst = soe_subst env
467 (subst_tv, tv') = substTyVarBndr subst bndr
468 (subst_cv, cv') = substCoVarBndr subst bndr
469
470 subst_opt_id_bndr :: SimpleOptEnv -> InId -> (SimpleOptEnv, OutId)
471 -- Nuke all fragile IdInfo, unfolding, and RULES;
472 -- it gets added back later by add_info
473 -- Rather like SimplEnv.substIdBndr
474 --
475 -- It's important to zap fragile OccInfo (which CoreSubst.substIdBndr
476 -- carefully does not do) because simplOptExpr invalidates it
477
478 subst_opt_id_bndr (SOE { soe_subst = subst, soe_inl = inl }) old_id
479 = (SOE { soe_subst = new_subst, soe_inl = new_inl }, new_id)
480 where
481 Subst in_scope id_subst tv_subst cv_subst = subst
482
483 id1 = uniqAway in_scope old_id
484 id2 = setIdType id1 (substTy subst (idType old_id))
485 new_id = zapFragileIdInfo id2
486 -- Zaps rules, unfolding, and fragile OccInfo
487 -- The unfolding and rules will get added back later, by add_info
488
489 new_in_scope = in_scope `extendInScopeSet` new_id
490
491 no_change = new_id == old_id
492
493 -- Extend the substitution if the unique has changed,
494 -- See the notes with substTyVarBndr for the delSubstEnv
495 new_id_subst
496 | no_change = delVarEnv id_subst old_id
497 | otherwise = extendVarEnv id_subst old_id (Var new_id)
498
499 new_subst = Subst new_in_scope new_id_subst tv_subst cv_subst
500 new_inl = delVarEnv inl old_id
501
502 ----------------------
503 add_info :: SimpleOptEnv -> InVar -> OutVar -> OutVar
504 add_info env old_bndr new_bndr
505 | isTyVar old_bndr = new_bndr
506 | otherwise = maybeModifyIdInfo mb_new_info new_bndr
507 where
508 subst = soe_subst env
509 mb_new_info = substIdInfo subst new_bndr (idInfo old_bndr)
510
511 simpleUnfoldingFun :: IdUnfoldingFun
512 simpleUnfoldingFun id
513 | isAlwaysActive (idInlineActivation id) = idUnfolding id
514 | otherwise = noUnfolding
515
516 wrapLet :: Maybe (Id,CoreExpr) -> CoreExpr -> CoreExpr
517 wrapLet Nothing body = body
518 wrapLet (Just (b,r)) body = Let (NonRec b r) body
519
520 ------------------
521 substVects :: Subst -> [CoreVect] -> [CoreVect]
522 substVects subst = map (substVect subst)
523
524 ------------------
525 substVect :: Subst -> CoreVect -> CoreVect
526 substVect subst (Vect v rhs) = Vect v (simpleOptExprWith subst rhs)
527 substVect _subst vd@(NoVect _) = vd
528 substVect _subst vd@(VectType _ _ _) = vd
529 substVect _subst vd@(VectClass _) = vd
530 substVect _subst vd@(VectInst _) = vd
531
532 {-
533 Note [Inline prag in simplOpt]
534 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
535 If there's an INLINE/NOINLINE pragma that restricts the phase in
536 which the binder can be inlined, we don't inline here; after all,
537 we don't know what phase we're in. Here's an example
538
539 foo :: Int -> Int -> Int
540 {-# INLINE foo #-}
541 foo m n = inner m
542 where
543 {-# INLINE [1] inner #-}
544 inner m = m+n
545
546 bar :: Int -> Int
547 bar n = foo n 1
548
549 When inlining 'foo' in 'bar' we want the let-binding for 'inner'
550 to remain visible until Phase 1
551
552 Note [Unfold compulsory unfoldings in LHSs]
553 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
554 When the user writes `RULES map coerce = coerce` as a rule, the rule
555 will only ever match if simpleOptExpr replaces coerce by its unfolding
556 on the LHS, because that is the core that the rule matching engine
557 will find. So do that for everything that has a compulsory
558 unfolding. Also see Note [Desugaring coerce as cast] in Desugar.
559
560 However, we don't want to inline 'seq', which happens to also have a
561 compulsory unfolding, so we only do this unfolding only for things
562 that are always-active. See Note [User-defined RULES for seq] in MkId.
563
564 Note [Getting the map/coerce RULE to work]
565 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
566 We wish to allow the "map/coerce" RULE to fire:
567
568 {-# RULES "map/coerce" map coerce = coerce #-}
569
570 The naive core produced for this is
571
572 forall a b (dict :: Coercible * a b).
573 map @a @b (coerce @a @b @dict) = coerce @[a] @[b] @dict'
574
575 where dict' :: Coercible [a] [b]
576 dict' = ...
577
578 This matches literal uses of `map coerce` in code, but that's not what we
579 want. We want it to match, say, `map MkAge` (where newtype Age = MkAge Int)
580 too. Some of this is addressed by compulsorily unfolding coerce on the LHS,
581 yielding
582
583 forall a b (dict :: Coercible * a b).
584 map @a @b (\(x :: a) -> case dict of
585 MkCoercible (co :: a ~R# b) -> x |> co) = ...
586
587 Getting better. But this isn't exactly what gets produced. This is because
588 Coercible essentially has ~R# as a superclass, and superclasses get eagerly
589 extracted during solving. So we get this:
590
591 forall a b (dict :: Coercible * a b).
592 case Coercible_SCSel @* @a @b dict of
593 _ [Dead] -> map @a @b (\(x :: a) -> case dict of
594 MkCoercible (co :: a ~R# b) -> x |> co) = ...
595
596 Unfortunately, this still abstracts over a Coercible dictionary. We really
597 want it to abstract over the ~R# evidence. So, we have Desugar.unfold_coerce,
598 which transforms the above to (see also Note [Desugaring coerce as cast] in
599 Desugar)
600
601 forall a b (co :: a ~R# b).
602 let dict = MkCoercible @* @a @b co in
603 case Coercible_SCSel @* @a @b dict of
604 _ [Dead] -> map @a @b (\(x :: a) -> case dict of
605 MkCoercible (co :: a ~R# b) -> x |> co) = let dict = ... in ...
606
607 Now, we need simpleOptExpr to fix this up. It does so by taking three
608 separate actions:
609 1. Inline certain non-recursive bindings. The choice whether to inline
610 is made in simple_bind_pair. Note the rather specific check for
611 MkCoercible in there.
612
613 2. Stripping case expressions like the Coercible_SCSel one.
614 See the `Case` case of simple_opt_expr's `go` function.
615
616 3. Look for case expressions that unpack something that was
617 just packed and inline them. This is also done in simple_opt_expr's
618 `go` function.
619
620 This is all a fair amount of special-purpose hackery, but it's for
621 a good cause. And it won't hurt other RULES and such that it comes across.
622
623
624 ************************************************************************
625 * *
626 Join points
627 * *
628 ************************************************************************
629 -}
630
631 -- | Returns Just (bndr,rhs) if the binding is a join point:
632 -- If it's a JoinId, just return it
633 -- If it's not yet a JoinId but is always tail-called,
634 -- make it into a JoinId and return it.
635 -- In the latter case, eta-expand the RHS if necessary, to make the
636 -- lambdas explicit, as is required for join points
637 --
638 -- Precondition: the InBndr has been occurrence-analysed,
639 -- so its OccInfo is valid
640 joinPointBinding_maybe :: InBndr -> InExpr -> Maybe (InBndr, InExpr)
641 joinPointBinding_maybe bndr rhs
642 | not (isId bndr)
643 = Nothing
644
645 | isJoinId bndr
646 = Just (bndr, rhs)
647
648 | AlwaysTailCalled join_arity <- tailCallInfo (idOccInfo bndr)
649 , not (bad_unfolding join_arity (idUnfolding bndr))
650 , (bndrs, body) <- etaExpandToJoinPoint join_arity rhs
651 = Just (bndr `asJoinId` join_arity, mkLams bndrs body)
652
653 | otherwise
654 = Nothing
655
656 where
657 -- bad_unfolding returns True if we should /not/ convert a non-join-id
658 -- into a join-id, even though it is AlwaysTailCalled
659 -- See Note [Join points and INLINE pragmas]
660 bad_unfolding join_arity (CoreUnfolding { uf_src = src, uf_tmpl = rhs })
661 = isStableSource src && join_arity > joinRhsArity rhs
662 bad_unfolding _ (DFunUnfolding {})
663 = True
664 bad_unfolding _ _
665 = False
666
667 joinPointBindings_maybe :: [(InBndr, InExpr)] -> Maybe [(InBndr, InExpr)]
668 joinPointBindings_maybe bndrs
669 = mapM (uncurry joinPointBinding_maybe) bndrs
670
671
672 {- Note [Join points and INLINE pragmas]
673 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
674 Consider
675 f x = let g = \x. not -- Arity 1
676 {-# INLINE g #-}
677 in case x of
678 A -> g True True
679 B -> g True False
680 C -> blah2
681
682 Here 'g' is always tail-called applied to 2 args, but the stable
683 unfolding captured by the INLINE pragma has arity 1. If we try to
684 convert g to be a join point, its unfolding will still have arity 1
685 (since it is stable, and we don't meddle with stable unfoldings), and
686 Lint will complain (see Note [Invariants on join points], (2a), in
687 CoreSyn. Trac #13413.
688
689 Moreover, since g is going to be inlined anyway, there is no benefit
690 from making it a join point.
691
692 If it is recursive, and uselessly marked INLINE, this will stop us
693 making it a join point, which is annoying. But occasionally
694 (notably in class methods; see Note [Instances and loop breakers] in
695 TcInstDcls) we mark recursive things as INLINE but the recursion
696 unravels; so ignoring INLINE pragmas on recursive things isn't good
697 either.
698
699
700 ************************************************************************
701 * *
702 exprIsConApp_maybe
703 * *
704 ************************************************************************
705
706 Note [exprIsConApp_maybe]
707 ~~~~~~~~~~~~~~~~~~~~~~~~~
708 exprIsConApp_maybe is a very important function. There are two principal
709 uses:
710 * case e of { .... }
711 * cls_op e, where cls_op is a class operation
712
713 In both cases you want to know if e is of form (C e1..en) where C is
714 a data constructor.
715
716 However e might not *look* as if
717
718
719 Note [exprIsConApp_maybe on literal strings]
720 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
721 See #9400 and #13317.
722
723 Conceptually, a string literal "abc" is just ('a':'b':'c':[]), but in Core
724 they are represented as unpackCString# "abc"# by MkCore.mkStringExprFS, or
725 unpackCStringUtf8# when the literal contains multi-byte UTF8 characters.
726
727 For optimizations we want to be able to treat it as a list, so they can be
728 decomposed when used in a case-statement. exprIsConApp_maybe detects those
729 calls to unpackCString# and returns:
730
731 Just (':', [Char], ['a', unpackCString# "bc"]).
732
733 We need to be careful about UTF8 strings here. ""# contains a ByteString, so
734 we must parse it back into a FastString to split off the first character.
735 That way we can treat unpackCString# and unpackCStringUtf8# in the same way.
736
737 We must also be caeful about
738 lvl = "foo"#
739 ...(unpackCString# lvl)...
740 to ensure that we see through the let-binding for 'lvl'. Hence the
741 (exprIsLiteral_maybe .. arg) in the guard before the call to
742 dealWithStringLiteral.
743
744 Note [Push coercions in exprIsConApp_maybe]
745 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
746 In Trac #13025 I found a case where we had
747 op (df @t1 @t2) -- op is a ClassOp
748 where
749 df = (/\a b. K e1 e2) |> g
750
751 To get this to come out we need to simplify on the fly
752 ((/\a b. K e1 e2) |> g) @t1 @t2
753
754 Hence the use of pushCoArgs.
755 -}
756
757 data ConCont = CC [CoreExpr] Coercion
758 -- Substitution already applied
759
760 -- | Returns @Just (dc, [t1..tk], [x1..xn])@ if the argument expression is
761 -- a *saturated* constructor application of the form @dc t1..tk x1 .. xn@,
762 -- where t1..tk are the *universally-quantified* type args of 'dc'
763 exprIsConApp_maybe :: InScopeEnv -> CoreExpr -> Maybe (DataCon, [Type], [CoreExpr])
764 exprIsConApp_maybe (in_scope, id_unf) expr
765 = go (Left in_scope) expr (CC [] (mkRepReflCo (exprType expr)))
766 where
767 go :: Either InScopeSet Subst
768 -- Left in-scope means "empty substitution"
769 -- Right subst means "apply this substitution to the CoreExpr"
770 -> CoreExpr -> ConCont
771 -> Maybe (DataCon, [Type], [CoreExpr])
772 go subst (Tick t expr) cont
773 | not (tickishIsCode t) = go subst expr cont
774 go subst (Cast expr co1) (CC args co2)
775 | Just (args', co1') <- pushCoArgs (subst_co subst co1) args
776 -- See Note [Push coercions in exprIsConApp_maybe]
777 = go subst expr (CC args' (co1' `mkTransCo` co2))
778 go subst (App fun arg) (CC args co)
779 = go subst fun (CC (subst_arg subst arg : args) co)
780 go subst (Lam var body) (CC (arg:args) co)
781 | exprIsTrivial arg -- Don't duplicate stuff!
782 = go (extend subst var arg) body (CC args co)
783 go (Right sub) (Var v) cont
784 = go (Left (substInScope sub))
785 (lookupIdSubst (text "exprIsConApp" <+> ppr expr) sub v)
786 cont
787
788 go (Left in_scope) (Var fun) cont@(CC args co)
789
790 | Just con <- isDataConWorkId_maybe fun
791 , count isValArg args == idArity fun
792 = pushCoDataCon con args co
793
794 -- Look through dictionary functions; see Note [Unfolding DFuns]
795 | DFunUnfolding { df_bndrs = bndrs, df_con = con, df_args = dfun_args } <- unfolding
796 , bndrs `equalLength` args -- See Note [DFun arity check]
797 , let subst = mkOpenSubst in_scope (bndrs `zip` args)
798 = pushCoDataCon con (map (substExpr (text "exprIsConApp1") subst) dfun_args) co
799
800 -- Look through unfoldings, but only arity-zero one;
801 -- if arity > 0 we are effectively inlining a function call,
802 -- and that is the business of callSiteInline.
803 -- In practice, without this test, most of the "hits" were
804 -- CPR'd workers getting inlined back into their wrappers,
805 | idArity fun == 0
806 , Just rhs <- expandUnfolding_maybe unfolding
807 , let in_scope' = extendInScopeSetSet in_scope (exprFreeVars rhs)
808 = go (Left in_scope') rhs cont
809
810 -- See Note [exprIsConApp_maybe on literal strings]
811 | (fun `hasKey` unpackCStringIdKey) ||
812 (fun `hasKey` unpackCStringUtf8IdKey)
813 , [arg] <- args
814 , Just (MachStr str) <- exprIsLiteral_maybe (in_scope, id_unf) arg
815 = dealWithStringLiteral fun str co
816 where
817 unfolding = id_unf fun
818
819 go _ _ _ = Nothing
820
821 ----------------------------
822 -- Operations on the (Either InScopeSet CoreSubst)
823 -- The Left case is wildly dominant
824 subst_co (Left {}) co = co
825 subst_co (Right s) co = CoreSubst.substCo s co
826
827 subst_arg (Left {}) e = e
828 subst_arg (Right s) e = substExpr (text "exprIsConApp2") s e
829
830 extend (Left in_scope) v e = Right (extendSubst (mkEmptySubst in_scope) v e)
831 extend (Right s) v e = Right (extendSubst s v e)
832
833
834 -- See Note [exprIsConApp_maybe on literal strings]
835 dealWithStringLiteral :: Var -> BS.ByteString -> Coercion
836 -> Maybe (DataCon, [Type], [CoreExpr])
837
838 -- This is not possible with user-supplied empty literals, MkCore.mkStringExprFS
839 -- turns those into [] automatically, but just in case something else in GHC
840 -- generates a string literal directly.
841 dealWithStringLiteral _ str co
842 | BS.null str
843 = pushCoDataCon nilDataCon [Type charTy] co
844
845 dealWithStringLiteral fun str co
846 = let strFS = mkFastStringByteString str
847
848 char = mkConApp charDataCon [mkCharLit (headFS strFS)]
849 charTail = fastStringToByteString (tailFS strFS)
850
851 -- In singleton strings, just add [] instead of unpackCstring# ""#.
852 rest = if BS.null charTail
853 then mkConApp nilDataCon [Type charTy]
854 else App (Var fun)
855 (Lit (MachStr charTail))
856
857 in pushCoDataCon consDataCon [Type charTy, char, rest] co
858
859 {-
860 Note [Unfolding DFuns]
861 ~~~~~~~~~~~~~~~~~~~~~~
862 DFuns look like
863
864 df :: forall a b. (Eq a, Eq b) -> Eq (a,b)
865 df a b d_a d_b = MkEqD (a,b) ($c1 a b d_a d_b)
866 ($c2 a b d_a d_b)
867
868 So to split it up we just need to apply the ops $c1, $c2 etc
869 to the very same args as the dfun. It takes a little more work
870 to compute the type arguments to the dictionary constructor.
871
872 Note [DFun arity check]
873 ~~~~~~~~~~~~~~~~~~~~~~~
874 Here we check that the total number of supplied arguments (inclding
875 type args) matches what the dfun is expecting. This may be *less*
876 than the ordinary arity of the dfun: see Note [DFun unfoldings] in CoreSyn
877 -}
878
879 exprIsLiteral_maybe :: InScopeEnv -> CoreExpr -> Maybe Literal
880 -- Same deal as exprIsConApp_maybe, but much simpler
881 -- Nevertheless we do need to look through unfoldings for
882 -- Integer and string literals, which are vigorously hoisted to top level
883 -- and not subsequently inlined
884 exprIsLiteral_maybe env@(_, id_unf) e
885 = case e of
886 Lit l -> Just l
887 Tick _ e' -> exprIsLiteral_maybe env e' -- dubious?
888 Var v | Just rhs <- expandUnfolding_maybe (id_unf v)
889 -> exprIsLiteral_maybe env rhs
890 _ -> Nothing
891
892 {-
893 Note [exprIsLambda_maybe]
894 ~~~~~~~~~~~~~~~~~~~~~~~~~~
895 exprIsLambda_maybe will, given an expression `e`, try to turn it into the form
896 `Lam v e'` (returned as `Just (v,e')`). Besides using lambdas, it looks through
897 casts (using the Push rule), and it unfolds function calls if the unfolding
898 has a greater arity than arguments are present.
899
900 Currently, it is used in Rules.match, and is required to make
901 "map coerce = coerce" match.
902 -}
903
904 exprIsLambda_maybe :: InScopeEnv -> CoreExpr
905 -> Maybe (Var, CoreExpr,[Tickish Id])
906 -- See Note [exprIsLambda_maybe]
907
908 -- The simple case: It is a lambda already
909 exprIsLambda_maybe _ (Lam x e)
910 = Just (x, e, [])
911
912 -- Still straightforward: Ticks that we can float out of the way
913 exprIsLambda_maybe (in_scope_set, id_unf) (Tick t e)
914 | tickishFloatable t
915 , Just (x, e, ts) <- exprIsLambda_maybe (in_scope_set, id_unf) e
916 = Just (x, e, t:ts)
917
918 -- Also possible: A casted lambda. Push the coercion inside
919 exprIsLambda_maybe (in_scope_set, id_unf) (Cast casted_e co)
920 | Just (x, e,ts) <- exprIsLambda_maybe (in_scope_set, id_unf) casted_e
921 -- Only do value lambdas.
922 -- this implies that x is not in scope in gamma (makes this code simpler)
923 , not (isTyVar x) && not (isCoVar x)
924 , ASSERT( not $ x `elemVarSet` tyCoVarsOfCo co) True
925 , Just (x',e') <- pushCoercionIntoLambda in_scope_set x e co
926 , let res = Just (x',e',ts)
927 = --pprTrace "exprIsLambda_maybe:Cast" (vcat [ppr casted_e,ppr co,ppr res)])
928 res
929
930 -- Another attempt: See if we find a partial unfolding
931 exprIsLambda_maybe (in_scope_set, id_unf) e
932 | (Var f, as, ts) <- collectArgsTicks tickishFloatable e
933 , idArity f > count isValArg as
934 -- Make sure there is hope to get a lambda
935 , Just rhs <- expandUnfolding_maybe (id_unf f)
936 -- Optimize, for beta-reduction
937 , let e' = simpleOptExprWith (mkEmptySubst in_scope_set) (rhs `mkApps` as)
938 -- Recurse, because of possible casts
939 , Just (x', e'', ts') <- exprIsLambda_maybe (in_scope_set, id_unf) e'
940 , let res = Just (x', e'', ts++ts')
941 = -- pprTrace "exprIsLambda_maybe:Unfold" (vcat [ppr e, ppr (x',e'')])
942 res
943
944 exprIsLambda_maybe _ _e
945 = -- pprTrace "exprIsLambda_maybe:Fail" (vcat [ppr _e])
946 Nothing
947
948
949 {- *********************************************************************
950 * *
951 The "push rules"
952 * *
953 ************************************************************************
954
955 Here we implement the "push rules" from FC papers:
956
957 * The push-argument rules, where we can move a coercion past an argument.
958 We have
959 (fun |> co) arg
960 and we want to transform it to
961 (fun arg') |> co'
962 for some suitable co' and tranformed arg'.
963
964 * The PushK rule for data constructors. We have
965 (K e1 .. en) |> co
966 and we want to tranform to
967 (K e1' .. en')
968 by pushing the coercion into the arguments
969 -}
970
971 pushCoArgs :: Coercion -> [CoreArg] -> Maybe ([CoreArg], Coercion)
972 pushCoArgs co [] = return ([], co)
973 pushCoArgs co (arg:args) = do { (arg', co1) <- pushCoArg co arg
974 ; (args', co2) <- pushCoArgs co1 args
975 ; return (arg':args', co2) }
976
977 pushCoArg :: Coercion -> CoreArg -> Maybe (CoreArg, Coercion)
978 -- We have (fun |> co) arg, and we want to transform it to
979 -- (fun arg) |> co
980 -- This may fail, e.g. if (fun :: N) where N is a newtype
981 -- C.f. simplCast in Simplify.hs
982 -- 'co' is always Representational
983
984 pushCoArg co (Type ty) = do { (ty', co') <- pushCoTyArg co ty
985 ; return (Type ty', co') }
986 pushCoArg co val_arg = do { (arg_co, co') <- pushCoValArg co
987 ; return (mkCast val_arg arg_co, co') }
988
989 pushCoTyArg :: Coercion -> Type -> Maybe (Type, Coercion)
990 -- We have (fun |> co) @ty
991 -- Push the coercion through to return
992 -- (fun @ty') |> co'
993 -- 'co' is always Representational
994 pushCoTyArg co ty
995 | tyL `eqType` tyR
996 = Just (ty, mkRepReflCo (piResultTy tyR ty))
997
998 | isForAllTy tyL
999 = ASSERT2( isForAllTy tyR, ppr co $$ ppr ty )
1000 Just (ty `mkCastTy` mkSymCo co1, co2)
1001
1002 | otherwise
1003 = Nothing
1004 where
1005 Pair tyL tyR = coercionKind co
1006 -- co :: tyL ~ tyR
1007 -- tyL = forall (a1 :: k1). ty1
1008 -- tyR = forall (a2 :: k2). ty2
1009
1010 co1 = mkNthCo 0 co
1011 -- co1 :: k1 ~ k2
1012 -- Note that NthCo can extract an equality between the kinds
1013 -- of the types related by a coercion between forall-types.
1014 -- See the NthCo case in CoreLint.
1015
1016 co2 = mkInstCo co (mkCoherenceLeftCo (mkNomReflCo ty) co1)
1017 -- co2 :: ty1[ (ty|>co1)/a1 ] ~ ty2[ ty/a2 ]
1018 -- Arg of mkInstCo is always nominal, hence mkNomReflCo
1019
1020 pushCoValArg :: Coercion -> Maybe (Coercion, Coercion)
1021 -- We have (fun |> co) arg
1022 -- Push the coercion through to return
1023 -- (fun (arg |> co_arg)) |> co_res
1024 -- 'co' is always Representational
1025 pushCoValArg co
1026 | tyL `eqType` tyR
1027 = Just (mkRepReflCo arg, mkRepReflCo res)
1028
1029 | isFunTy tyL
1030 , (co1, co2) <- decomposeFunCo co
1031 -- If co :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
1032 -- then co1 :: tyL1 ~ tyR1
1033 -- co2 :: tyL2 ~ tyR2
1034 = ASSERT2( isFunTy tyR, ppr co $$ ppr arg )
1035 Just (mkSymCo co1, co2)
1036
1037 | otherwise
1038 = Nothing
1039 where
1040 (arg, res) = splitFunTy tyR
1041 Pair tyL tyR = coercionKind co
1042
1043 pushCoercionIntoLambda
1044 :: InScopeSet -> Var -> CoreExpr -> Coercion -> Maybe (Var, CoreExpr)
1045 -- This implements the Push rule from the paper on coercions
1046 -- (\x. e) |> co
1047 -- ===>
1048 -- (\x'. e |> co')
1049 pushCoercionIntoLambda in_scope x e co
1050 | ASSERT(not (isTyVar x) && not (isCoVar x)) True
1051 , Pair s1s2 t1t2 <- coercionKind co
1052 , Just (_s1,_s2) <- splitFunTy_maybe s1s2
1053 , Just (t1,_t2) <- splitFunTy_maybe t1t2
1054 = let (co1, co2) = decomposeFunCo co
1055 -- Should we optimize the coercions here?
1056 -- Otherwise they might not match too well
1057 x' = x `setIdType` t1
1058 in_scope' = in_scope `extendInScopeSet` x'
1059 subst = extendIdSubst (mkEmptySubst in_scope')
1060 x
1061 (mkCast (Var x') co1)
1062 in Just (x', substExpr (text "pushCoercionIntoLambda") subst e `mkCast` co2)
1063 | otherwise
1064 = pprTrace "exprIsLambda_maybe: Unexpected lambda in case" (ppr (Lam x e))
1065 Nothing
1066
1067 pushCoDataCon :: DataCon -> [CoreExpr] -> Coercion
1068 -> Maybe (DataCon
1069 , [Type] -- Universal type args
1070 , [CoreExpr]) -- All other args incl existentials
1071 -- Implement the KPush reduction rule as described in "Down with kinds"
1072 -- The transformation applies iff we have
1073 -- (C e1 ... en) `cast` co
1074 -- where co :: (T t1 .. tn) ~ to_ty
1075 -- The left-hand one must be a T, because exprIsConApp returned True
1076 -- but the right-hand one might not be. (Though it usually will.)
1077 pushCoDataCon dc dc_args co
1078 | isReflCo co || from_ty `eqType` to_ty -- try cheap test first
1079 , let (univ_ty_args, rest_args) = splitAtList (dataConUnivTyVars dc) dc_args
1080 = Just (dc, map exprToType univ_ty_args, rest_args)
1081
1082 | Just (to_tc, to_tc_arg_tys) <- splitTyConApp_maybe to_ty
1083 , to_tc == dataConTyCon dc
1084 -- These two tests can fail; we might see
1085 -- (C x y) `cast` (g :: T a ~ S [a]),
1086 -- where S is a type function. In fact, exprIsConApp
1087 -- will probably not be called in such circumstances,
1088 -- but there's nothing wrong with it
1089
1090 = let
1091 tc_arity = tyConArity to_tc
1092 dc_univ_tyvars = dataConUnivTyVars dc
1093 dc_ex_tyvars = dataConExTyVars dc
1094 arg_tys = dataConRepArgTys dc
1095
1096 non_univ_args = dropList dc_univ_tyvars dc_args
1097 (ex_args, val_args) = splitAtList dc_ex_tyvars non_univ_args
1098
1099 -- Make the "Psi" from the paper
1100 omegas = decomposeCo tc_arity co
1101 (psi_subst, to_ex_arg_tys)
1102 = liftCoSubstWithEx Representational
1103 dc_univ_tyvars
1104 omegas
1105 dc_ex_tyvars
1106 (map exprToType ex_args)
1107
1108 -- Cast the value arguments (which include dictionaries)
1109 new_val_args = zipWith cast_arg arg_tys val_args
1110 cast_arg arg_ty arg = mkCast arg (psi_subst arg_ty)
1111
1112 to_ex_args = map Type to_ex_arg_tys
1113
1114 dump_doc = vcat [ppr dc, ppr dc_univ_tyvars, ppr dc_ex_tyvars,
1115 ppr arg_tys, ppr dc_args,
1116 ppr ex_args, ppr val_args, ppr co, ppr from_ty, ppr to_ty, ppr to_tc ]
1117 in
1118 ASSERT2( eqType from_ty (mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args)), dump_doc )
1119 ASSERT2( equalLength val_args arg_tys, dump_doc )
1120 Just (dc, to_tc_arg_tys, to_ex_args ++ new_val_args)
1121
1122 | otherwise
1123 = Nothing
1124
1125 where
1126 Pair from_ty to_ty = coercionKind co
1127
1128 collectBindersPushingCo :: CoreExpr -> ([Var], CoreExpr)
1129 -- Collect lambda binders, pushing coercions inside if possible
1130 -- E.g. (\x.e) |> g g :: <Int> -> blah
1131 -- = (\x. e |> Nth 1 g)
1132 --
1133 -- That is,
1134 --
1135 -- collectBindersPushingCo ((\x.e) |> g) === ([x], e |> Nth 1 g)
1136 collectBindersPushingCo e
1137 = go [] e
1138 where
1139 -- Peel off lambdas until we hit a cast.
1140 go :: [Var] -> CoreExpr -> ([Var], CoreExpr)
1141 -- The accumulator is in reverse order
1142 go bs (Lam b e) = go (b:bs) e
1143 go bs (Cast e co) = go_c bs e co
1144 go bs e = (reverse bs, e)
1145
1146 -- We are in a cast; peel off casts until we hit a lambda.
1147 go_c :: [Var] -> CoreExpr -> Coercion -> ([Var], CoreExpr)
1148 -- (go_c bs e c) is same as (go bs e (e |> c))
1149 go_c bs (Cast e co1) co2 = go_c bs e (co1 `mkTransCo` co2)
1150 go_c bs (Lam b e) co = go_lam bs b e co
1151 go_c bs e co = (reverse bs, mkCast e co)
1152
1153 -- We are in a lambda under a cast; peel off lambdas and build a
1154 -- new coercion for the body.
1155 go_lam :: [Var] -> Var -> CoreExpr -> Coercion -> ([Var], CoreExpr)
1156 -- (go_lam bs b e c) is same as (go_c bs (\b.e) c)
1157 go_lam bs b e co
1158 | isTyVar b
1159 , let Pair tyL tyR = coercionKind co
1160 , ASSERT( isForAllTy tyL )
1161 isForAllTy tyR
1162 , isReflCo (mkNthCo 0 co) -- See Note [collectBindersPushingCo]
1163 = go_c (b:bs) e (mkInstCo co (mkNomReflCo (mkTyVarTy b)))
1164
1165 | isId b
1166 , let Pair tyL tyR = coercionKind co
1167 , ASSERT( isFunTy tyL) isFunTy tyR
1168 , (co_arg, co_res) <- decomposeFunCo co
1169 , isReflCo co_arg -- See Note [collectBindersPushingCo]
1170 = go_c (b:bs) e co_res
1171
1172 | otherwise = (reverse bs, mkCast (Lam b e) co)
1173
1174 {- Note [collectBindersPushingCo]
1175 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1176 We just look for coercions of form
1177 <type> -> blah
1178 (and similarly for foralls) to keep this function simple. We could do
1179 more elaborate stuff, but it'd involve substitution etc.
1180 -}