2 % (c) The University of Glasgow 2006
6 module OptCoercion ( optCoercion ) where
8 #include "HsVersions.h"
11 import Type hiding( substTyVarBndr, substTy, extendTvSubst )
16 import StaticFlags ( opt_NoOptCoercion )
19 import Maybes( allMaybes )
23 %************************************************************************
27 %************************************************************************
29 Note [Subtle shadowing in coercions]
30 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
31 Supose we optimising a coercion
32 optCoercion (forall (co_X5:t1~t2). ...co_B1...)
33 The co_X5 is a wild-card; the bound variable of a coercion for-all
34 should never appear in the body of the forall. Indeed we often
36 optCoercion ( (t1~t2) => ...co_B1... )
38 Just because it's a wild-card doesn't mean we are free to choose
39 whatever variable we like. For example it'd be wrong for optCoercion
41 forall (co_B1:t1~t2). ...co_B1...
42 because now the co_B1 (which is really free) has been captured, and
43 subsequent substitutions will go wrong. That's why we can't use
44 mkCoPredTy in the ForAll case, where this note appears.
47 optCoercion :: CvSubst -> Coercion -> NormalCo
48 -- ^ optCoercion applies a substitution to a coercion,
49 -- *and* optimises it to reduce its size
51 | opt_NoOptCoercion = substCo env co
52 | otherwise = opt_co env False co
54 type NormalCo = Coercion
56 -- * The substitution has been fully applied
57 -- * For trans coercions (co1 `trans` co2)
58 -- co1 is not a trans, and neither co1 nor co2 is identity
59 -- * If the coercion is the identity, it has no CoVars of CoTyCons in it (just types)
61 type NormalNonIdCo = NormalCo -- Extra invariant: not the identity
63 opt_co, opt_co' :: CvSubst
64 -> Bool -- True <=> return (sym co)
70 = pprTrace "opt_co {" (ppr sym <+> ppr co $$ ppr env) $
72 pprTrace "opt_co done }" (ppr co1) $
73 (WARN( not same_co_kind, ppr co <+> dcolon <+> pprEqPred (Pair s1 t1)
74 $$ ppr co1 <+> dcolon <+> pprEqPred (Pair s2 t2) )
75 WARN( not (coreEqCoercion co1 simple_result),
76 (text "env=" <+> ppr env) $$
77 (text "input=" <+> ppr co) $$
78 (text "simple=" <+> ppr simple_result) $$
79 (text "opt=" <+> ppr co1) )
82 co1 = opt_co' env sym co
83 same_co_kind = s1 `eqType` s2 && t1 `eqType` t2
84 Pair s t = coercionKind (substCo env co)
87 Pair s2 t2 = coercionKind co1
89 simple_result | sym = mkSymCo (substCo env co)
90 | otherwise = substCo env co
93 opt_co' env _ (Refl ty) = Refl (substTy env ty)
94 opt_co' env sym (SymCo co) = opt_co env (not sym) co
95 opt_co' env sym (TyConAppCo tc cos) = mkTyConAppCo tc (map (opt_co env sym) cos)
96 opt_co' env sym (AppCo co1 co2) = mkAppCo (opt_co env sym co1) (opt_co env sym co2)
97 opt_co' env sym (ForAllCo tv co) = case substTyVarBndr env tv of
98 (env', tv') -> mkForAllCo tv' (opt_co env' sym co)
99 -- Use the "mk" functions to check for nested Refls
101 opt_co' env sym (CoVarCo cv)
102 | Just co <- lookupCoVar env cv
103 = opt_co (zapCvSubstEnv env) sym co
105 | Just cv1 <- lookupInScope (getCvInScope env) cv
106 = ASSERT( isCoVar cv1 ) wrapSym sym (CoVarCo cv1)
107 -- cv1 might have a substituted kind!
109 | otherwise = WARN( True, ptext (sLit "opt_co: not in scope:") <+> ppr cv $$ ppr env)
111 wrapSym sym (CoVarCo cv)
113 opt_co' env sym (AxiomInstCo con cos)
114 -- Do *not* push sym inside top-level axioms
115 -- e.g. if g is a top-level axiom
117 -- then (sym (g ty)) /= g (sym ty) !!
118 = wrapSym sym $ AxiomInstCo con (map (opt_co env False) cos)
119 -- Note that the_co does *not* have sym pushed into it
121 opt_co' env sym (UnsafeCo ty1 ty2)
122 | ty1' `eqType` ty2' = Refl ty1'
123 | sym = mkUnsafeCo ty2' ty1'
124 | otherwise = mkUnsafeCo ty1' ty2'
126 ty1' = substTy env ty1
127 ty2' = substTy env ty2
129 opt_co' env sym (TransCo co1 co2)
130 | sym = opt_trans opt_co2 opt_co1 -- sym (g `o` h) = sym h `o` sym g
131 | otherwise = opt_trans opt_co1 opt_co2
133 opt_co1 = opt_co env sym co1
134 opt_co2 = opt_co env sym co2
136 opt_co' env sym (NthCo n co)
137 | TyConAppCo tc cos <- co'
138 , isDecomposableTyCon tc -- Not synonym families
139 = ASSERT( n < length cos )
144 co' = opt_co env sym co
146 opt_co' env sym (InstCo co ty)
147 -- See if the first arg is already a forall
148 -- ...then we can just extend the current substitution
149 | Just (tv, co_body) <- splitForAllCo_maybe co
150 = opt_co (extendTvSubst env tv ty') sym co_body
152 -- See if it is a forall after optimization
153 | Just (tv, co'_body) <- splitForAllCo_maybe co'
154 = substCoWithTy tv ty' co'_body -- An inefficient one-variable substitution
156 | otherwise = InstCo co' ty'
159 co' = opt_co env sym co
163 opt_transList :: [NormalCo] -> [NormalCo] -> [NormalCo]
164 opt_transList = zipWith opt_trans
166 opt_trans :: NormalCo -> NormalCo -> NormalCo
169 | otherwise = opt_trans1 co1 co2
171 opt_trans1 :: NormalNonIdCo -> NormalCo -> NormalCo
172 -- First arg is not the identity
175 | otherwise = opt_trans2 co1 co2
177 opt_trans2 :: NormalNonIdCo -> NormalNonIdCo -> NormalCo
178 -- Neither arg is the identity
179 opt_trans2 (TransCo co1a co1b) co2
180 -- Don't know whether the sub-coercions are the identity
181 = opt_trans co1a (opt_trans co1b co2)
184 | Just co <- opt_trans_rule co1 co2
187 opt_trans2 co1 (TransCo co2a co2b)
188 | Just co1_2a <- opt_trans_rule co1 co2a
191 else opt_trans1 co1_2a co2b
197 -- Optimize coercions with a top-level use of transitivity.
198 opt_trans_rule :: NormalNonIdCo -> NormalNonIdCo -> Maybe NormalCo
200 -- push transitivity down through matching top-level constructors.
201 opt_trans_rule in_co1@(TyConAppCo tc1 cos1) in_co2@(TyConAppCo tc2 cos2)
202 | tc1 == tc2 = fireTransRule "PushTyConApp" in_co1 in_co2 $
203 TyConAppCo tc1 (opt_transList cos1 cos2)
205 -- push transitivity through matching destructors
206 opt_trans_rule in_co1@(NthCo d1 co1) in_co2@(NthCo d2 co2)
208 , co1 `compatible_co` co2
209 = fireTransRule "PushNth" in_co1 in_co2 $
210 mkNthCo d1 (opt_trans co1 co2)
212 -- Push transitivity inside instantiation
213 opt_trans_rule in_co1@(InstCo co1 ty1) in_co2@(InstCo co2 ty2)
215 , co1 `compatible_co` co2
216 = fireTransRule "TrPushInst" in_co1 in_co2 $
217 mkInstCo (opt_trans co1 co2) ty1
219 -- Push transitivity inside apply
220 opt_trans_rule in_co1@(AppCo co1a co1b) in_co2@(AppCo co2a co2b)
221 = fireTransRule "TrPushApp" in_co1 in_co2 $
222 mkAppCo (opt_trans co1a co2a) (opt_trans co1b co2b)
224 opt_trans_rule co1@(TyConAppCo tc cos1) co2
225 | Just cos2 <- etaTyConAppCo_maybe tc co2
226 = ASSERT( length cos1 == length cos2 )
227 fireTransRule "EtaCompL" co1 co2 $
228 TyConAppCo tc (zipWith opt_trans cos1 cos2)
230 opt_trans_rule co1 co2@(TyConAppCo tc cos2)
231 | Just cos1 <- etaTyConAppCo_maybe tc co1
232 = ASSERT( length cos1 == length cos2 )
233 fireTransRule "EtaCompR" co1 co2 $
234 TyConAppCo tc (zipWith opt_trans cos1 cos2)
236 -- Push transitivity inside forall
237 opt_trans_rule co1 co2
238 | Just (tv1,r1) <- splitForAllCo_maybe co1
239 , Just (tv2,r2) <- etaForAllCo_maybe co2
240 , let r2' = substCoWithTy tv2 (mkTyVarTy tv1) r2
241 = fireTransRule "EtaAllL" co1 co2 $
242 mkForAllCo tv1 (opt_trans2 r1 r2')
244 | Just (tv2,r2) <- splitForAllCo_maybe co2
245 , Just (tv1,r1) <- etaForAllCo_maybe co1
246 , let r1' = substCoWithTy tv1 (mkTyVarTy tv2) r1
247 = fireTransRule "EtaAllR" co1 co2 $
248 mkForAllCo tv1 (opt_trans2 r1' r2)
250 -- Push transitivity inside axioms
251 opt_trans_rule co1 co2
253 -- TrPushAxR/TrPushSymAxR
254 | Just (sym, con, cos1) <- co1_is_axiom_maybe
255 , Just cos2 <- matchAxiom sym con co2
256 = fireTransRule "TrPushAxR" co1 co2 $
258 then SymCo $ AxiomInstCo con (opt_transList (map mkSymCo cos2) cos1)
259 else AxiomInstCo con (opt_transList cos1 cos2)
261 -- TrPushAxL/TrPushSymAxL
262 | Just (sym, con, cos2) <- co2_is_axiom_maybe
263 , Just cos1 <- matchAxiom (not sym) con co1
264 = fireTransRule "TrPushAxL" co1 co2 $
266 then SymCo $ AxiomInstCo con (opt_transList cos2 (map mkSymCo cos1))
267 else AxiomInstCo con (opt_transList cos1 cos2)
269 -- TrPushAxSym/TrPushSymAx
270 | Just (sym1, con1, cos1) <- co1_is_axiom_maybe
271 , Just (sym2, con2, cos2) <- co2_is_axiom_maybe
274 , let qtvs = co_ax_tvs con1
277 pivot_tvs = exactTyVarsOfType (if sym2 then rhs else lhs)
278 , all (`elemVarSet` pivot_tvs) qtvs
279 = fireTransRule "TrPushAxSym" co1 co2 $
281 then liftCoSubstWith qtvs (opt_transList cos1 (map mkSymCo cos2)) lhs -- TrPushAxSym
282 else liftCoSubstWith qtvs (opt_transList (map mkSymCo cos1) cos2) rhs -- TrPushSymAx
284 co1_is_axiom_maybe = isAxiom_maybe co1
285 co2_is_axiom_maybe = isAxiom_maybe co2
287 opt_trans_rule co1 co2 -- Identity rule
288 | Pair ty1 _ <- coercionKind co1
289 , Pair _ ty2 <- coercionKind co2
291 = fireTransRule "RedTypeDirRefl" co1 co2 $
294 opt_trans_rule _ _ = Nothing
296 fireTransRule :: String -> Coercion -> Coercion -> Coercion -> Maybe Coercion
297 fireTransRule _rule _co1 _co2 res
298 = -- pprTrace ("Trans rule fired: " ++ _rule) (vcat [ppr _co1, ppr _co2, ppr res]) $
302 wrapSym :: Bool -> Coercion -> Coercion
303 wrapSym sym co | sym = SymCo co
307 isAxiom_maybe :: Coercion -> Maybe (Bool, CoAxiom, [Coercion])
308 isAxiom_maybe (SymCo co)
309 | Just (sym, con, cos) <- isAxiom_maybe co
310 = Just (not sym, con, cos)
311 isAxiom_maybe (AxiomInstCo con cos)
312 = Just (False, con, cos)
313 isAxiom_maybe _ = Nothing
315 matchAxiom :: Bool -- True = match LHS, False = match RHS
316 -> CoAxiom -> Coercion -> Maybe [Coercion]
317 -- If we succeed in matching, then *all the quantified type variables are bound*
318 -- E.g. if tvs = [a,b], lhs/rhs = [b], we'll fail
319 matchAxiom sym (CoAxiom { co_ax_tvs = qtvs, co_ax_lhs = lhs, co_ax_rhs = rhs }) co
320 = case liftCoMatch (mkVarSet qtvs) (if sym then lhs else rhs) co of
322 Just subst -> allMaybes (map (liftCoSubstTyVar subst) qtvs)
325 compatible_co :: Coercion -> Coercion -> Bool
326 -- Check whether (co1 . co2) will be well-kinded
327 compatible_co co1 co2
330 Pair _ x1 = coercionKind co1
331 Pair x2 _ = coercionKind co2
334 etaForAllCo_maybe :: Coercion -> Maybe (TyVar, Coercion)
335 -- Try to make the coercion be of form (forall tv. co)
337 | Just (tv, r) <- splitForAllCo_maybe co
340 | Pair ty1 ty2 <- coercionKind co
341 , Just (tv1, _) <- splitForAllTy_maybe ty1
342 , Just (tv2, _) <- splitForAllTy_maybe ty2
343 , tyVarKind tv1 `eqKind` tyVarKind tv2
344 = Just (tv1, mkInstCo co (mkTyVarTy tv1))
349 etaTyConAppCo_maybe :: TyCon -> Coercion -> Maybe [Coercion]
350 -- If possible, split a coercion
351 -- g :: T s1 .. sn ~ T t1 .. tn
352 -- into [ Nth 0 g :: s1~t1, ..., Nth (n-1) g :: sn~tn ]
353 etaTyConAppCo_maybe tc (TyConAppCo tc2 cos2)
354 = ASSERT( tc == tc2 ) Just cos2
356 etaTyConAppCo_maybe tc co
357 | isDecomposableTyCon tc
358 , Pair ty1 ty2 <- coercionKind co
359 , Just (tc1, tys1) <- splitTyConApp_maybe ty1
360 , Just (tc2, tys2) <- splitTyConApp_maybe ty2
362 , let n = length tys1
363 = ASSERT( tc == tc1 )
364 ASSERT( n == length tys2 )
365 Just (decomposeCo n co)
366 -- NB: n might be <> tyConArity tc
367 -- e.g. data family T a :: * -> *
368 -- g :: T a b ~ T c d