compiler/types/OptCoercion.lhs

   1 %
   2 % (c) The University of Glasgow 2006
   3 %
   4
   5 \begin{code}
   6 module OptCoercion ( optCoercion ) where
   7
   8 #include "HsVersions.h"
   9
  10 import Coercion
  11 import Type hiding( substTyVarBndr, substTy, extendTvSubst )
  12 import TyCon
  13 import Var
  14 import VarSet
  15 import VarEnv
  16 import StaticFlags      ( opt_NoOptCoercion )
  17 import Outputable
  18 import Pair
  19 import Maybes( allMaybes )
  20 import FastString
  21 \end{code}
  22
  23 %************************************************************************
  24 %*                                                                      *
  25                  Optimising coercions
  26 %*                                                                      *
  27 %************************************************************************
  28
  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
  35 write it like this
  36     optCoercion ( (t1~t2) => ...co_B1... )
  37
  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
  40 to return
  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.
  45
  46 \begin{code}
  47 optCoercion :: CvSubst -> Coercion -> NormalCo
  48 -- ^ optCoercion applies a substitution to a coercion,
  49 --   *and* optimises it to reduce its size
  50 optCoercion env co
  51   | opt_NoOptCoercion = substCo env co
  52   | otherwise         = opt_co env False co
  53
  54 type NormalCo = Coercion
  55   -- Invariants:
  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)
  60
  61 type NormalNonIdCo = NormalCo  -- Extra invariant: not the identity
  62
  63 opt_co, opt_co' :: CvSubst
  64                 -> Bool        -- True <=> return (sym co)
  65                 -> Coercion
  66                 -> NormalCo
  67 opt_co = opt_co'
  68 {-
  69 opt_co env sym co
  70  = pprTrace "opt_co {" (ppr sym <+> ppr co $$ ppr env) $
  71    co1 `seq`
  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) )
  80    co1)
  81  where
  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)
  85    (s1,t1) | sym = (t,s)
  86            | otherwise = (s,t)
  87    Pair s2 t2 = coercionKind co1
  88
  89    simple_result | sym = mkSymCo (substCo env co)
  90                  | otherwise = substCo env co
  91 -}
  92
  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
 100
 101 opt_co' env sym (CoVarCo cv)
 102   | Just co <- lookupCoVar env cv
 103   = opt_co (zapCvSubstEnv env) sym co
 104
 105   | Just cv1 <- lookupInScope (getCvInScope env) cv
 106   = ASSERT( isCoVar cv1 ) wrapSym sym (CoVarCo cv1)
 107                 -- cv1 might have a substituted kind!
 108
 109   | otherwise = WARN( True, ptext (sLit "opt_co: not in scope:") <+> ppr cv $$ ppr env)
 110                 ASSERT( isCoVar cv )
 111                 wrapSym sym (CoVarCo cv)
 112
 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
 116     --   g a : f a ~ a
 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
 120
 121 opt_co' env sym (UnsafeCo ty1 ty2)
 122   | ty1' `eqType` ty2' = Refl ty1'
 123   | sym                = mkUnsafeCo ty2' ty1'
 124   | otherwise          = mkUnsafeCo ty1' ty2'
 125   where
 126     ty1' = substTy env ty1
 127     ty2' = substTy env ty2
 128
 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
 132   where
 133     opt_co1 = opt_co env sym co1
 134     opt_co2 = opt_co env sym co2
 135
 136 opt_co' env sym (NthCo n co)
 137   | TyConAppCo tc cos <- co'
 138   , isDecomposableTyCon tc              -- Not synonym families
 139   = ASSERT( n < length cos )
 140     cos !! n
 141   | otherwise
 142   = NthCo n co'
 143   where
 144     co' = opt_co env sym co
 145
 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
 151
 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
 155
 156   | otherwise = InstCo co' ty'
 157
 158   where
 159     co' = opt_co env sym co
 160     ty' = substTy env ty
 161
 162 -------------
 163 opt_transList :: [NormalCo] -> [NormalCo] -> [NormalCo]
 164 opt_transList = zipWith opt_trans
 165
 166 opt_trans :: NormalCo -> NormalCo -> NormalCo
 167 opt_trans co1 co2
 168   | isReflCo co1 = co2
 169   | otherwise    = opt_trans1 co1 co2
 170
 171 opt_trans1 :: NormalNonIdCo -> NormalCo -> NormalCo
 172 -- First arg is not the identity
 173 opt_trans1 co1 co2
 174   | isReflCo co2 = co1
 175   | otherwise    = opt_trans2 co1 co2
 176
 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)
 182
 183 opt_trans2 co1 co2
 184   | Just co <- opt_trans_rule co1 co2
 185   = co
 186
 187 opt_trans2 co1 (TransCo co2a co2b)
 188   | Just co1_2a <- opt_trans_rule co1 co2a
 189   = if isReflCo co1_2a
 190     then co2b
 191     else opt_trans1 co1_2a co2b
 192
 193 opt_trans2 co1 co2
 194   = mkTransCo co1 co2
 195
 196 ------
 197 -- Optimize coercions with a top-level use of transitivity.
 198 opt_trans_rule :: NormalNonIdCo -> NormalNonIdCo -> Maybe NormalCo
 199
 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)
 204
 205 -- push transitivity through matching destructors
 206 opt_trans_rule in_co1@(NthCo d1 co1) in_co2@(NthCo d2 co2)
 207   | d1 == d2
 208   , co1 `compatible_co` co2
 209   = fireTransRule "PushNth" in_co1 in_co2 $
 210     mkNthCo d1 (opt_trans co1 co2)
 211
 212 -- Push transitivity inside instantiation
 213 opt_trans_rule in_co1@(InstCo co1 ty1) in_co2@(InstCo co2 ty2)
 214   | ty1 `eqType` ty2
 215   , co1 `compatible_co` co2
 216   = fireTransRule "TrPushInst" in_co1 in_co2 $
 217     mkInstCo (opt_trans co1 co2) ty1
 218
 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)
 223
 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)
 229
 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)
 235
 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')
 243
 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)
 249
 250 -- Push transitivity inside axioms
 251 opt_trans_rule co1 co2
 252
 253   -- TrPushAxR/TrPushSymAxR
 254   | Just (sym, con, cos1) <- co1_is_axiom_maybe
 255   , Just cos2 <- matchAxiom sym con co2
 256   = fireTransRule "TrPushAxR" co1 co2 $
 257     if sym
 258     then SymCo $ AxiomInstCo con (opt_transList (map mkSymCo cos2) cos1)
 259     else         AxiomInstCo con (opt_transList cos1 cos2)
 260
 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 $
 265     if sym
 266     then SymCo $ AxiomInstCo con (opt_transList cos2 (map mkSymCo cos1))
 267     else         AxiomInstCo con (opt_transList cos1 cos2)
 268
 269   -- TrPushAxSym/TrPushSymAx
 270   | Just (sym1, con1, cos1) <- co1_is_axiom_maybe
 271   , Just (sym2, con2, cos2) <- co2_is_axiom_maybe
 272   , con1 == con2
 273   , sym1 == not sym2
 274   , let qtvs = co_ax_tvs con1
 275         lhs  = co_ax_lhs con1
 276         rhs  = co_ax_rhs con1
 277         pivot_tvs = exactTyVarsOfType (if sym2 then rhs else lhs)
 278   , all (`elemVarSet` pivot_tvs) qtvs
 279   = fireTransRule "TrPushAxSym" co1 co2 $
 280     if sym2
 281     then liftCoSubstWith qtvs (opt_transList cos1 (map mkSymCo cos2)) lhs  -- TrPushAxSym
 282     else liftCoSubstWith qtvs (opt_transList (map mkSymCo cos1) cos2) rhs  -- TrPushSymAx
 283   where
 284     co1_is_axiom_maybe = isAxiom_maybe co1
 285     co2_is_axiom_maybe = isAxiom_maybe co2
 286
 287 opt_trans_rule co1 co2  -- Identity rule
 288   | Pair ty1 _ <- coercionKind co1
 289   , Pair _ ty2 <- coercionKind co2
 290   , ty1 `eqType` ty2
 291   = fireTransRule "RedTypeDirRefl" co1 co2 $
 292     Refl ty2
 293
 294 opt_trans_rule _ _ = Nothing
 295
 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]) $
 299     Just res
 300
 301 -----------
 302 wrapSym :: Bool -> Coercion -> Coercion
 303 wrapSym sym co | sym       = SymCo co
 304                | otherwise = co
 305
 306 -----------
 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
 314
 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
 321       Nothing    -> Nothing
 322       Just subst -> allMaybes (map (liftCoSubstTyVar subst) qtvs)
 323
 324 -------------
 325 compatible_co :: Coercion -> Coercion -> Bool
 326 -- Check whether (co1 . co2) will be well-kinded
 327 compatible_co co1 co2
 328   = x1 `eqType` x2
 329   where
 330     Pair _ x1 = coercionKind co1
 331     Pair x2 _ = coercionKind co2
 332
 333 -------------
 334 etaForAllCo_maybe :: Coercion -> Maybe (TyVar, Coercion)
 335 -- Try to make the coercion be of form (forall tv. co)
 336 etaForAllCo_maybe co
 337   | Just (tv, r) <- splitForAllCo_maybe co
 338   = Just (tv, r)
 339
 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))
 345
 346   | otherwise
 347   = Nothing
 348
 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
 355
 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
 361   , tc1 == tc2
 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
 369
 370   | otherwise
 371   = Nothing
 372 \end{code}