Substantial improvements to coercion optimisation
[ghc-hetmet.git] / compiler / types / OptCoercion.lhs
diff --git a/compiler/types/OptCoercion.lhs b/compiler/types/OptCoercion.lhs
new file mode 100644 (file)
index 0000000..43de7d6
--- /dev/null
@@ -0,0 +1,403 @@
+%\r
+% (c) The University of Glasgow 2006\r
+%\r
+\r
+\begin{code}\r
+{-# OPTIONS_GHC -w #-}\r
+module OptCoercion (\r
+       optCoercion\r
+   ) where \r
+\r
+#include "HsVersions.h"\r
+\r
+import Unify   ( tcMatchTy )\r
+import Coercion\r
+import Type\r
+import TypeRep\r
+import TyCon\r
+import Var\r
+import VarSet\r
+import PrelNames\r
+import Util\r
+import Outputable\r
+\end{code}\r
+\r
+%************************************************************************\r
+%*                                                                      *\r
+                 Optimising coercions                                                                  \r
+%*                                                                      *\r
+%************************************************************************\r
+\r
+\begin{code}\r
+optCoercion :: TvSubst -> Coercion -> NormalCo\r
+-- ^ optCoercion applies a substitution to a coercion, \r
+--   *and* optimises it to reduce its size\r
+optCoercion env co = opt_co env False co\r
+\r
+type NormalCo = Coercion\r
+  -- Invariants: \r
+  --  * The substitution has been fully applied\r
+  --  * For trans coercions (co1 `trans` co2)\r
+  --       co1 is not a trans, and neither co1 nor co2 is identity\r
+  --  * If the coercion is the identity, it has no CoVars of CoTyCons in it (just types)\r
+\r
+type NormalNonIdCo = NormalCo  -- Extra invariant: not the identity\r
+\r
+opt_co, opt_co' :: TvSubst\r
+                       -> Bool        -- True <=> return (sym co)\r
+                       -> Coercion\r
+                       -> NormalCo     \r
+opt_co = opt_co'\r
+-- opt_co sym co = pprTrace "opt_co {" (ppr sym <+> ppr co) $\r
+--                     co1 `seq` \r
+--                pprTrace "opt_co done }" (ppr co1) \r
+--               WARN( not same_co_kind, ppr co  <+> dcolon <+> pprEqPred (s1,t1) \r
+--                                     $$ ppr co1 <+> dcolon <+> pprEqPred (s2,t2) )\r
+--                co1\r
+--  where\r
+--    co1 = opt_co' sym co\r
+--    same_co_kind = s1 `coreEqType` s2 && t1 `coreEqType` t2\r
+--    (s,t) = coercionKind co\r
+--    (s1,t1) | sym = (t,s)\r
+--            | otherwise = (s,t)\r
+--    (s2,t2) = coercionKind co1\r
+\r
+opt_co' env sym (AppTy ty1 ty2)          = mkAppTy (opt_co env sym ty1) (opt_co env sym ty2)\r
+opt_co' env sym (FunTy ty1 ty2)          = FunTy (opt_co env sym ty1) (opt_co env sym ty2)\r
+opt_co' env sym (PredTy (ClassP cls tys)) = PredTy (ClassP cls (map (opt_co env sym) tys))\r
+opt_co' env sym (PredTy (IParam n ty))    = PredTy (IParam n (opt_co env sym ty))\r
+opt_co' _   _   co@(PredTy (EqPred {}))   = pprPanic "optCoercion" (ppr co)\r
+\r
+opt_co' env sym co@(TyVarTy tv)\r
+  | Just ty <- lookupTyVar env tv = opt_co' (zapTvSubstEnv env) sym ty\r
+  | not (isCoVar tv)     = co   -- Identity; does not mention a CoVar\r
+  | ty1 `coreEqType` ty2 = ty1 -- Identity; ..ditto..\r
+  | not sym              = co\r
+  | otherwise            = mkSymCoercion co\r
+  where\r
+    (ty1,ty2) = coVarKind tv\r
+\r
+opt_co' env sym (ForAllTy tv cor) \r
+  | isCoVar tv = mkCoPredTy (opt_co env sym co1) (opt_co env sym co2) (opt_co env sym cor)\r
+  | otherwise  = case substTyVarBndr env tv of\r
+                   (env', tv') -> ForAllTy tv' (opt_co env' sym cor)\r
+  where\r
+    (co1,co2) = coVarKind tv\r
+\r
+opt_co' env sym (TyConApp tc cos)\r
+  | Just (arity, desc) <- isCoercionTyCon_maybe tc\r
+  = mkAppTys (opt_co_tc_app env sym tc desc (take arity cos))\r
+             (map (opt_co env sym) (drop arity cos))\r
+  | otherwise\r
+  = TyConApp tc (map (opt_co env sym) cos)\r
+\r
+--------\r
+opt_co_tc_app :: TvSubst -> Bool -> TyCon -> CoTyConDesc -> [Coercion] -> NormalCo\r
+-- Used for CoercionTyCons only\r
+-- Arguments are *not* already simplified/substituted\r
+opt_co_tc_app env sym tc desc cos\r
+  = case desc of\r
+      CoAxiom {} -- Do *not* push sym inside top-level axioms\r
+                -- e.g. if g is a top-level axiom\r
+                --   g a : F a ~ a\r
+                -- Then (sym (g ty)) /= g (sym ty) !!\r
+        | sym       -> mkSymCoercion the_co  \r
+        | otherwise -> the_co\r
+        where\r
+           the_co = TyConApp tc (map (opt_co env False) cos)\r
+           -- Note that the_co does *not* have sym pushed into it\r
+    \r
+      CoTrans \r
+        | sym       -> opt_trans opt_co2 opt_co1   -- sym (g `o` h) = sym h `o` sym g\r
+        | otherwise -> opt_trans opt_co1 opt_co2\r
+\r
+      CoUnsafe\r
+        | sym       -> TyConApp tc [opt_co2,opt_co1]\r
+        | otherwise -> TyConApp tc [opt_co1,opt_co2]\r
+\r
+      CoSym   -> opt_co env (not sym) co1\r
+      CoLeft  -> opt_lr fst\r
+      CoRight -> opt_lr snd\r
+      CoCsel1 -> opt_csel fstOf3\r
+      CoCsel2 -> opt_csel sndOf3\r
+      CoCselR -> opt_csel thirdOf3\r
+\r
+      CoInst        -- See if the first arg is already a forall\r
+                   -- ...then we can just extend the current substitution\r
+        | Just (tv, co1_body) <- splitForAllTy_maybe co1\r
+        -> opt_co (extendTvSubst env tv ty') sym co1_body\r
+\r
+                    -- See if is *now* a forall\r
+        | Just (tv, opt_co1_body) <- splitForAllTy_maybe opt_co1\r
+        -> substTyWith [tv] [ty'] opt_co1_body -- An inefficient one-variable substitution\r
+\r
+        | otherwise\r
+        -> TyConApp tc [opt_co1, ty']\r
+        where\r
+          ty' = substTy env co2\r
+\r
+  where\r
+    (co1 : cos1) = cos\r
+    (co2 : _)    = cos1\r
+\r
+       -- These opt_cos have the sym pushed into them\r
+    opt_co1 = opt_co env sym co1\r
+    opt_co2 = opt_co env sym co2\r
+\r
+    the_unary_opt_co = TyConApp tc [opt_co1]\r
+\r
+    opt_lr   sel = case splitAppTy_maybe opt_co1 of\r
+                     Nothing -> the_unary_opt_co \r
+                     Just lr -> sel lr\r
+    opt_csel sel = case splitCoPredTy_maybe opt_co1 of\r
+                     Nothing -> the_unary_opt_co \r
+                     Just lr -> sel lr\r
+\r
+-------------\r
+opt_transL :: [NormalCo] -> [NormalCo] -> [NormalCo]\r
+opt_transL = zipWith opt_trans\r
+\r
+opt_trans :: NormalCo -> NormalCo -> NormalCo\r
+opt_trans co1 co2\r
+  | isIdNormCo co1 = co2\r
+  | otherwise      = opt_trans1 co1 co2\r
+\r
+opt_trans1 :: NormalNonIdCo -> NormalCo -> NormalCo\r
+-- First arg is not the identity\r
+opt_trans1 co1 co2\r
+  | isIdNormCo co2 = co1\r
+  | otherwise      = opt_trans2 co1 co2\r
+\r
+opt_trans2 :: NormalNonIdCo -> NormalNonIdCo -> NormalCo\r
+-- Neither arg is the identity\r
+opt_trans2 (TyConApp tc [co1a,co1b]) co2\r
+  | tc `hasKey` transCoercionTyConKey\r
+  = opt_trans1 co1a (opt_trans2 co1b co2)\r
+\r
+opt_trans2 co1 co2 \r
+  | Just co <- opt_trans_rule co1 co2\r
+  = co\r
+\r
+opt_trans2 co1 (TyConApp tc [co2a,co2b])\r
+  | tc `hasKey` transCoercionTyConKey\r
+  , Just co1_2a <- opt_trans_rule co1 co2a\r
+  = if isIdNormCo co1_2a\r
+    then co2b\r
+    else opt_trans2 co1_2a co2b\r
+\r
+opt_trans2 co1 co2\r
+  = mkTransCoercion co1 co2\r
+\r
+------\r
+opt_trans_rule :: NormalNonIdCo -> NormalNonIdCo -> Maybe NormalCo\r
+opt_trans_rule (TyConApp tc1 args1) (TyConApp tc2 args2)\r
+  | tc1 == tc2\r
+  = case isCoercionTyCon_maybe tc1 of\r
+      Nothing \r
+        -> Just (TyConApp tc1 (opt_transL args1 args2))\r
+      Just (arity, desc) \r
+        | arity == length args1\r
+        -> opt_trans_rule_equal_tc desc args1 args2\r
+        | otherwise\r
+        -> case opt_trans_rule_equal_tc desc \r
+                         (take arity args1) \r
+                         (take arity args2) of\r
+              Just co -> Just $ mkAppTys co $ \r
+                         opt_transL (drop arity args1) (drop arity args2)\r
+             Nothing -> Nothing \r
\r
+-- Push transitivity inside apply\r
+opt_trans_rule co1 co2\r
+  | Just (co1a, co1b) <- splitAppTy_maybe co1\r
+  , Just (co2a, co2b) <- etaApp_maybe co2\r
+  = Just (mkAppTy (opt_trans co1a co2a) (opt_trans co1b co2b))\r
+\r
+  | Just (co2a, co2b) <- splitAppTy_maybe co2\r
+  , Just (co1a, co1b) <- etaApp_maybe co1\r
+  = Just (mkAppTy (opt_trans co1a co2a) (opt_trans co1b co2b))\r
+\r
+-- Push transitivity inside (s~t)=>r\r
+opt_trans_rule co1 co2\r
+  | Just (s1,t1,r1) <- splitCoPredTy_maybe co1\r
+  , Just (s2,t2,r2) <- etaCoPred_maybe co2\r
+  = Just (mkCoPredTy (opt_trans s1 s2) (opt_trans t1 t2) (opt_trans r1 r2))\r
+\r
+  | Just (s2,t2,r2) <- splitCoPredTy_maybe co2\r
+  , Just (s1,t1,r1) <- etaCoPred_maybe co1\r
+  = Just (mkCoPredTy (opt_trans s1 s2) (opt_trans t1 t2) (opt_trans r1 r2))\r
+\r
+-- Push transitivity inside forall\r
+opt_trans_rule co1 co2\r
+  | Just (tv1,r1) <- splitTypeForAll_maybe co1\r
+  , Just (tv2,r2) <- etaForAll_maybe co2\r
+  , let r2' = substTyWith [tv2] [TyVarTy tv1] r2\r
+  = Just (ForAllTy tv1 (opt_trans2 r1 r2'))\r
+\r
+  | Just (tv2,r2) <- splitTypeForAll_maybe co2\r
+  , Just (tv1,r1) <- etaForAll_maybe co1\r
+  , let r1' = substTyWith [tv1] [TyVarTy tv2] r1\r
+  = Just (ForAllTy tv1 (opt_trans2 r1' r2))\r
+\r
+opt_trans_rule co1 co2\r
+{-     Omitting for now, because unsound\r
+  | Just (sym1, (ax_tc1, ax1_args, ax_tvs, ax_lhs, ax_rhs)) <- co1_is_axiom_maybe\r
+  , Just (sym2, (ax_tc2, ax2_args, _, _, _)) <- co2_is_axiom_maybe\r
+  , ax_tc1 == ax_tc2\r
+  , sym1 /= sym2\r
+  = Just $\r
+    if sym1 \r
+    then substTyWith ax_tvs (opt_transL (map mkSymCoercion ax1_args) ax2_args) ax_rhs\r
+    else substTyWith ax_tvs (opt_transL ax1_args (map mkSymCoercion ax2_args)) ax_lhs\r
+-}\r
+\r
+  | Just (sym, (ax_tc, ax_args, ax_tvs, ax_lhs, _)) <- co1_is_axiom_maybe\r
+  , Just cos <- matchesAxiomLhs ax_tvs ax_lhs co2\r
+  = Just $ \r
+    if sym \r
+    then mkSymCoercion $ TyConApp ax_tc (opt_transL (map mkSymCoercion cos) ax_args)\r
+    else                 TyConApp ax_tc (opt_transL ax_args cos)\r
+\r
+  | Just (sym, (ax_tc, ax_args, ax_tvs, ax_lhs, _)) <- isAxiom_maybe co2\r
+  , Just cos <- matchesAxiomLhs ax_tvs ax_lhs co1\r
+  = Just $ \r
+    if sym \r
+    then mkSymCoercion $ TyConApp ax_tc (opt_transL ax_args (map mkSymCoercion cos))\r
+    else                 TyConApp ax_tc (opt_transL cos ax_args)\r
+  where\r
+    co1_is_axiom_maybe = isAxiom_maybe co1\r
+    co2_is_axiom_maybe = isAxiom_maybe co2\r
+\r
+opt_trans_rule co1 co2 -- Identity rule\r
+  | (ty1,_) <- coercionKind co1\r
+  , (_,ty2) <- coercionKind co2\r
+  , ty1 `coreEqType` ty2\r
+  = Just ty2\r
+\r
+opt_trans_rule _ _ = Nothing\r
+\r
+-----------  \r
+isAxiom_maybe :: Coercion -> Maybe (Bool, (TyCon, [Coercion], [TyVar], Type, Type))\r
+isAxiom_maybe co\r
+  | Just (tc, args) <- splitTyConApp_maybe co\r
+  , Just (_, desc)  <- isCoercionTyCon_maybe tc\r
+  = case desc of\r
+      CoAxiom { co_ax_tvs = tvs, co_ax_lhs = lhs, co_ax_rhs = rhs } \r
+            -> Just (False, (tc, args, tvs, lhs, rhs))\r
+      CoSym | (arg1:_) <- args  \r
+            -> case isAxiom_maybe arg1 of\r
+                 Nothing           -> Nothing\r
+                 Just (sym, stuff) -> Just (not sym, stuff)\r
+      _ -> Nothing\r
+  | otherwise\r
+  = Nothing\r
+\r
+matchesAxiomLhs :: [TyVar] -> Type -> Type -> Maybe [Type]\r
+matchesAxiomLhs tvs ty_tmpl ty \r
+  = case tcMatchTy (mkVarSet tvs) ty_tmpl ty of\r
+      Nothing    -> Nothing\r
+      Just subst -> Just (map (substTyVar subst) tvs)\r
+\r
+-----------  \r
+opt_trans_rule_equal_tc :: CoTyConDesc -> [Coercion] -> [Coercion] -> Maybe Coercion\r
+-- Rules for Coercion TyCons only\r
+\r
+-- Push transitivity inside instantiation\r
+opt_trans_rule_equal_tc desc [co1,ty1] [co2,ty2]\r
+  | CoInst <- desc\r
+  , ty1 `coreEqType` ty2\r
+  , co1 `compatible_co` co2\r
+  = Just (mkInstCoercion (opt_trans2 co1 co2) ty1) \r
+\r
+opt_trans_rule_equal_tc desc [co1] [co2]\r
+  | CoLeft  <- desc, is_compat = Just (mkLeftCoercion res_co)\r
+  | CoRight <- desc, is_compat = Just (mkRightCoercion res_co)\r
+  | CoCsel1 <- desc, is_compat = Just (mkCsel1Coercion res_co)\r
+  | CoCsel2 <- desc, is_compat = Just (mkCsel2Coercion res_co)\r
+  | CoCselR <- desc, is_compat = Just (mkCselRCoercion res_co)\r
+  where\r
+    is_compat = co1 `compatible_co` co2\r
+    res_co    = opt_trans2 co1 co2\r
+\r
+opt_trans_rule_equal_tc _ _ _ = Nothing\r
+\r
+-------------\r
+compatible_co :: Coercion -> Coercion -> Bool\r
+-- Check whether (co1 . co2) will be well-kinded\r
+compatible_co co1 co2\r
+  = x1 `coreEqType` x2         \r
+  where\r
+    (_,x1) = coercionKind co1\r
+    (x2,_) = coercionKind co2\r
+\r
+-------------\r
+etaForAll_maybe :: Coercion -> Maybe (TyVar, Coercion)\r
+-- Try to make the coercion be of form (forall tv. co)\r
+etaForAll_maybe co\r
+  | Just (tv, r) <- splitForAllTy_maybe co\r
+  , not (isCoVar tv)   -- Check it is a *type* forall, not a (t1~t2)=>co\r
+  = Just (tv, r)\r
+\r
+  | (ty1,ty2) <- coercionKind co\r
+  , Just (tv1, _) <- splitTypeForAll_maybe ty1\r
+  , Just (tv2, _) <- splitTypeForAll_maybe ty2\r
+  , tyVarKind tv1 `eqKind` tyVarKind tv2\r
+  = Just (tv1, mkInstCoercion co (mkTyVarTy tv1))\r
+\r
+  | otherwise\r
+  = Nothing\r
+\r
+etaCoPred_maybe :: Coercion -> Maybe (Coercion, Coercion, Coercion)\r
+etaCoPred_maybe co \r
+  | Just (s,t,r) <- splitCoPredTy_maybe co\r
+  = Just (s,t,r)\r
+  \r
+  --  co :: (s1~t1)=>r1 ~ (s2~t2)=>r2\r
+  | (ty1,ty2) <- coercionKind co       -- We know ty1,ty2 have same kind\r
+  , Just (s1,_,_) <- splitCoPredTy_maybe ty1\r
+  , Just (s2,_,_) <- splitCoPredTy_maybe ty2\r
+  , typeKind s1 `eqKind` typeKind s2   -- t1,t2 have same kinds\r
+  = Just (mkCsel1Coercion co, mkCsel2Coercion co, mkCselRCoercion co)\r
+  \r
+  | otherwise\r
+  = Nothing\r
+\r
+etaApp_maybe :: Coercion -> Maybe (Coercion, Coercion)\r
+etaApp_maybe co\r
+  | Just (co1, co2) <- splitAppTy_maybe co\r
+  = Just (co1, co2)\r
+\r
+  | (ty1,ty2) <- coercionKind co\r
+  , Just (ty1a, _) <- splitAppTy_maybe ty1\r
+  , Just (ty2a, _) <- splitAppTy_maybe ty2\r
+  , typeKind ty1a `eqKind` typeKind ty2a\r
+  = Just (mkLeftCoercion co, mkRightCoercion co)\r
+\r
+  | otherwise\r
+  = Nothing\r
+\r
+-------------\r
+splitTypeForAll_maybe :: Type -> Maybe (TyVar, Type)\r
+-- Returns Just only for a *type* forall, not a (t1~t2)=>co\r
+splitTypeForAll_maybe ty\r
+  | Just (tv, rty) <- splitForAllTy_maybe ty\r
+  , not (isCoVar tv)\r
+  = Just (tv, rty)\r
+\r
+  | otherwise\r
+  = Nothing\r
+\r
+-------------\r
+isIdNormCo :: NormalCo -> Bool\r
+-- Cheap identity test: look for coercions with no coercion variables at all\r
+-- So it'll return False for (sym g `trans` g)\r
+isIdNormCo ty = go ty\r
+  where\r
+    go (TyVarTy tv)           = not (isCoVar tv)\r
+    go (AppTy t1 t2)          = go t1 && go t2\r
+    go (FunTy t1 t2)          = go t1 && go t2\r
+    go (ForAllTy tv ty)        = go (tyVarKind tv) && go ty\r
+    go (TyConApp tc tys)       = not (isCoercionTyCon tc) && all go tys\r
+    go (PredTy (IParam _ ty))  = go ty\r
+    go (PredTy (ClassP _ tys)) = all go tys\r
+    go (PredTy (EqPred t1 t2)) = go t1 && go t2\r
+\end{code}  \r