Directed coercionswip/amg/dcoercion

This patch introduces a slimmer version of coercions, directed coercions,
which store fewer types within them. This more compact representation
considerably speeds up programs which involve many type family
reductions, as the coercion size no longer grows quadratically in
the number of reduction steps.

-------------------------
Metric Decrease:
    LargeRecord
    T12227
    T12545
    T13386
    T3064
    T5030
    T8095
    T9872a
    T9872b
    T9872b_defer
    T9872c
    T9872d
Metric Increase:
    CoOpt_Singletons
    T18223
    T9872a
    T9872b
    T9872c
    T9872d
-------------------------

1 files changed, 135 insertions, 45 deletions

diff --git a/compiler/GHC/Core/TyCo/Subst.hs b/compiler/GHC/Core/TyCo/Subst.hs
index 4224bd127b..ee25e16f75 100644
--- a/compiler/GHC/Core/TyCo/Subst.hs
+++ b/compiler/GHC/Core/TyCo/Subst.hs
@@ -6,6 +6,7 @@ Type and Coercion - friends' interface
 
 
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE GADTs #-}
 
 -- | Substitution into types and coercions.
 module GHC.Core.TyCo.Subst
@@ -35,17 +36,19 @@ module GHC.Core.TyCo.Subst
         substTyUnchecked, substTysUnchecked, substScaledTysUnchecked, substThetaUnchecked,
         substTyWithUnchecked, substScaledTyUnchecked,
         substCoUnchecked, substCoWithUnchecked,
+        substDCoUnchecked,
         substTyWithInScope,
         substTys, substScaledTys, substTheta,
         lookupTyVar,
         substCo, substCos, substCoVar, substCoVars, lookupCoVar,
+        substDCo,
         cloneTyVarBndr, cloneTyVarBndrs,
         substVarBndr, substVarBndrs,
         substTyVarBndr, substTyVarBndrs,
         substCoVarBndr,
         substTyVar, substTyVars, substTyVarToTyVar,
         substTyCoVars,
-        substTyCoBndr, substForAllCoBndr,
+        substTyCoBndr, substForAllDCoBndr, substForAllCoBndr,
         substVarBndrUsing, substForAllCoBndrUsing,
         checkValidSubst, isValidTCvSubst,
   ) where
@@ -61,7 +64,11 @@ import {-# SOURCE #-} GHC.Core.Coercion
    , mkAxiomInstCo, mkAppCo, mkGReflCo
    , mkInstCo, mkLRCo, mkTyConAppCo
    , mkCoercionType
-   , coercionKind, coercionLKind, coVarKindsTypesRole )
+   , mkTyConAppDCo
+   , mkAppDCo, mkForAllDCo, mkReflDCo, mkTransDCo
+   , mkGReflRightDCo, mkGReflLeftDCo
+   , mkHydrateDCo, mkDehydrateCo, mkUnivDCo
+   , coercionKind, coercionLKind, coVarKindsTypesRole)
 import {-# SOURCE #-} GHC.Core.TyCo.Ppr ( pprTyVar )
 import {-# SOURCE #-} GHC.Core.Ppr ( )
 import {-# SOURCE #-} GHC.Core ( CoreExpr )
@@ -682,8 +689,8 @@ isValidTCvSubst (Subst in_scope _ tenv cenv) =
 
 -- | This checks if the substitution satisfies the invariant from
 -- Note [The substitution invariant].
-checkValidSubst :: HasDebugCallStack => Subst -> [Type] -> [Coercion] -> a -> a
-checkValidSubst subst@(Subst in_scope _ tenv cenv) tys cos a
+checkValidSubst :: HasDebugCallStack => Subst -> [Type] -> [Coercion] -> [DCoercion] -> a -> a
+checkValidSubst subst@(Subst in_scope _ tenv cenv) tys cos dcos a
   = assertPpr (isValidTCvSubst subst)
               (text "in_scope" <+> ppr in_scope $$
                text "tenv" <+> ppr tenv $$
@@ -691,13 +698,15 @@ checkValidSubst subst@(Subst in_scope _ tenv cenv) tys cos a
                text "cenv" <+> ppr cenv $$
                text "cenvFVs" <+> ppr (shallowTyCoVarsOfCoVarEnv cenv) $$
                text "tys" <+> ppr tys $$
-               text "cos" <+> ppr cos) $
+               text "cos" <+> ppr cos $$
+               text "dcos" <+> ppr dcos) $
     assertPpr tysCosFVsInScope
               (text "in_scope" <+> ppr in_scope $$
                text "tenv" <+> ppr tenv $$
                text "cenv" <+> ppr cenv $$
                text "tys" <+> ppr tys $$
                text "cos" <+> ppr cos $$
+               text "dcos" <+> ppr dcos $$
                text "needInScope" <+> ppr needInScope)
     a
   where
@@ -705,7 +714,8 @@ checkValidSubst subst@(Subst in_scope _ tenv cenv) tys cos a
     -- It's OK to use nonDetKeysUFM here, because we only use this list to
     -- remove some elements from a set
   needInScope = (shallowTyCoVarsOfTypes tys `unionVarSet`
-                 shallowTyCoVarsOfCos cos)
+                 shallowTyCoVarsOfCos cos `unionVarSet`
+                 shallowTyCoVarsOfDCos dcos)
                 `delListFromUniqSet_Directly` substDomain
   tysCosFVsInScope = needInScope `varSetInScope` in_scope
 
@@ -716,7 +726,7 @@ checkValidSubst subst@(Subst in_scope _ tenv cenv) tys cos a
 substTy :: HasDebugCallStack => Subst -> Type  -> Type
 substTy subst ty
   | isEmptyTCvSubst    subst = ty
-  | otherwise             = checkValidSubst subst [ty] [] $
+  | otherwise             = checkValidSubst subst [ty] [] [] $
                             subst_ty subst ty
 
 -- | Substitute within a 'Type' disabling the sanity checks.
@@ -741,12 +751,12 @@ substScaledTyUnchecked subst scaled_ty = mapScaledType (substTyUnchecked subst)
 substTys :: HasDebugCallStack => Subst -> [Type] -> [Type]
 substTys subst tys
   | isEmptyTCvSubst subst = tys
-  | otherwise = checkValidSubst subst tys [] $ map (subst_ty subst) tys
+  | otherwise = checkValidSubst subst tys [] [] $ map (subst_ty subst) tys
 
 substScaledTys :: HasDebugCallStack => Subst -> [Scaled Type] -> [Scaled Type]
 substScaledTys subst scaled_tys
   | isEmptyTCvSubst subst = scaled_tys
-  | otherwise = checkValidSubst subst (map scaledMult scaled_tys ++ map scaledThing scaled_tys) [] $
+  | otherwise = checkValidSubst subst (map scaledMult scaled_tys ++ map scaledThing scaled_tys) [] [] $
                 map (mapScaledType (subst_ty subst)) scaled_tys
 
 -- | Substitute within several 'Type's disabling the sanity checks.
@@ -846,13 +856,31 @@ lookupTyVar (Subst _ _ tenv _) tv
   = assert (isTyVar tv )
     lookupVarEnv tenv tv
 
+-- | Substitute within a 'DCoercion'
+-- The substitution has to satisfy the invariants described in
+-- Note [The substitution invariant].
+substDCo :: HasDebugCallStack => Subst -> DCoercion -> DCoercion
+substDCo subst dco
+  | isEmptyTCvSubst subst = dco
+  | otherwise = checkValidSubst subst [] [] [dco] $ subst_dco subst dco
+
+-- | Substitute within a 'DCoercion' disabling sanity checks.
+-- The problems that the sanity checks in substCo catch are described in
+-- Note [The substitution invariant].
+-- The goal of #11371 is to migrate all the calls of substDCoUnchecked to
+-- substDCo and remove this function. Please don't use in new code.
+substDCoUnchecked :: Subst -> DCoercion -> DCoercion
+substDCoUnchecked subst co
+  | isEmptyTCvSubst subst = co
+  | otherwise = subst_dco subst co
+
 -- | Substitute within a 'Coercion'
 -- The substitution has to satisfy the invariants described in
 -- Note [The substitution invariant].
 substCo :: HasDebugCallStack => Subst -> Coercion -> Coercion
 substCo subst co
   | isEmptyTCvSubst subst = co
-  | otherwise = checkValidSubst subst [] [co] $ subst_co subst co
+  | otherwise = checkValidSubst subst [] [co] [] $ subst_co subst co
 
 -- | Substitute within a 'Coercion' disabling sanity checks.
 -- The problems that the sanity checks in substCo catch are described in
@@ -870,18 +898,23 @@ substCoUnchecked subst co
 substCos :: HasDebugCallStack => Subst -> [Coercion] -> [Coercion]
 substCos subst cos
   | isEmptyTCvSubst subst = cos
-  | otherwise = checkValidSubst subst [] cos $ map (subst_co subst) cos
+  | otherwise = checkValidSubst subst [] cos [] $ map (subst_co subst) cos
 
 subst_co :: Subst -> Coercion -> Coercion
-subst_co subst co
-  = go co
+subst_co = fst . subst_co_dco
+
+subst_dco :: Subst -> DCoercion -> DCoercion
+subst_dco = snd . subst_co_dco
+
+subst_co_dco :: Subst -> (Coercion -> Coercion, DCoercion -> DCoercion)
+subst_co_dco subst = (go, go_dco)
   where
     go_ty :: Type -> Type
     go_ty = subst_ty subst
 
     go_mco :: MCoercion -> MCoercion
     go_mco MRefl    = MRefl
-    go_mco (MCo co) = MCo (go co)
+    go_mco (MCo co) = MCo $! go co
 
     go :: Coercion -> Coercion
     go (Refl ty)             = mkNomReflCo $! (go_ty ty)
@@ -896,7 +929,10 @@ subst_co subst co
     go (FunCo r afl afr w co1 co2)   = ((mkFunCo2 r afl afr $! go w) $! go co1) $! go co2
     go (CoVarCo cv)          = substCoVar subst cv
     go (AxiomInstCo con ind cos) = mkAxiomInstCo con ind $! map go cos
-    go (UnivCo p r t1 t2)    = (((mkUnivCo $! go_prov p) $! r) $!
+    go (HydrateDCo r ty dco rty) = (((mkHydrateDCo $! r) $! go_ty ty) $! go_dco dco) $! go_ty rty
+      -- Here we can either substitute the RHS or recompute it from the rest of the information.
+
+    go (UnivCo p r t1 t2)    = (((mkUnivCo $! go_prov go p) $! r) $!
                                 (go_ty t1)) $! (go_ty t2)
     go (SymCo co)            = mkSymCo $! (go co)
     go (TransCo co1 co2)     = (mkTransCo $! (go co1)) $! (go co2)
@@ -909,10 +945,30 @@ subst_co subst co
                                 in cs1 `seqList` AxiomRuleCo c cs1
     go (HoleCo h)            = HoleCo $! go_hole h
 
-    go_prov (PhantomProv kco)    = PhantomProv (go kco)
-    go_prov (ProofIrrelProv kco) = ProofIrrelProv (go kco)
-    go_prov p@(PluginProv _)     = p
-    go_prov p@(CorePrepProv _)   = p
+    go_dco :: DCoercion -> DCoercion
+    go_dco ReflDCo                = mkReflDCo
+    go_dco (GReflRightDCo co)     = mkGReflRightDCo $! go co
+    go_dco (GReflLeftDCo  co)     = mkGReflLeftDCo  $! go co
+    go_dco (TyConAppDCo args)     = let args' = map go_dco args
+                                    in  args' `seqList` mkTyConAppDCo args'
+    go_dco (AppDCo co arg)        = (mkAppDCo $! go_dco co) $! go_dco arg
+    go_dco (CoVarDCo cv)          = mkDehydrateCo $! substCoVar subst cv
+    go_dco dco@AxiomInstDCo{}     = dco
+    go_dco dco@StepsDCo{}         = dco
+    go_dco (TransDCo co1 co2)     = (mkTransDCo $! go_dco co1) $! go_dco co2
+    go_dco (DehydrateCo co)       = mkDehydrateCo $! go co
+    go_dco (ForAllDCo tv kind_dco co)
+      = case substForAllDCoBndrUnchecked subst tv kind_dco of
+         (subst', tv', kind_dco') ->
+          ((mkForAllDCo $! tv') $! kind_dco') $! subst_dco subst' co
+    go_dco (UnivDCo prov rhs)     = (mkUnivDCo (go_prov go_dco prov)) $! go_ty rhs
+    go_dco (SubDCo dco)           = SubDCo $ go_dco dco
+
+    go_prov :: (co -> co) -> UnivCoProvenance co -> UnivCoProvenance co
+    go_prov do_subst (PhantomProv kco)    = PhantomProv $! do_subst kco
+    go_prov do_subst (ProofIrrelProv kco) = ProofIrrelProv $! do_subst kco
+    go_prov _        p@(PluginProv _)     = p
+    go_prov _        p@(CorePrepProv _)   = p
 
     -- See Note [Substituting in a coercion hole]
     go_hole h@(CoercionHole { ch_co_var = cv })
@@ -921,7 +977,12 @@ subst_co subst co
 substForAllCoBndr :: Subst -> TyCoVar -> KindCoercion
                   -> (Subst, TyCoVar, Coercion)
 substForAllCoBndr subst
-  = substForAllCoBndrUsing False (substCo subst) subst
+  = substForAllCoBndrUsing Co False (substTy subst) (substCo subst) subst
+
+substForAllDCoBndr :: Subst -> TyCoVar -> KindDCoercion
+                  -> (Subst, TyCoVar, DCoercion)
+substForAllDCoBndr subst
+  = substForAllCoBndrUsing DCo False (substTy subst) (substDCo subst) subst
 
 -- | Like 'substForAllCoBndr', but disables sanity checks.
 -- The problems that the sanity checks in substCo catch are described in
@@ -931,65 +992,94 @@ substForAllCoBndr subst
 substForAllCoBndrUnchecked :: Subst -> TyCoVar -> KindCoercion
                            -> (Subst, TyCoVar, Coercion)
 substForAllCoBndrUnchecked subst
-  = substForAllCoBndrUsing False (substCoUnchecked subst) subst
+  = substForAllCoBndrUsing Co False (substTyUnchecked subst) (substCoUnchecked subst) subst
+
+substForAllDCoBndrUnchecked :: Subst -> TyCoVar -> KindDCoercion
+                           -> (Subst, TyCoVar, DCoercion)
+substForAllDCoBndrUnchecked subst
+  = substForAllCoBndrUsing DCo False (substTyUnchecked subst) (substDCoUnchecked subst) subst
+
 
 -- See Note [Sym and ForAllCo]
-substForAllCoBndrUsing :: Bool  -- apply sym to binder?
-                       -> (Coercion -> Coercion)  -- transformation to kind co
-                       -> Subst -> TyCoVar -> KindCoercion
-                       -> (Subst, TyCoVar, KindCoercion)
-substForAllCoBndrUsing sym sco subst old_var
-  | isTyVar old_var = substForAllCoTyVarBndrUsing sym sco subst old_var
-  | otherwise       = substForAllCoCoVarBndrUsing sym sco subst old_var
-
-substForAllCoTyVarBndrUsing :: Bool  -- apply sym to binder?
-                            -> (Coercion -> Coercion)  -- transformation to kind co
-                            -> Subst -> TyVar -> KindCoercion
-                            -> (Subst, TyVar, KindCoercion)
-substForAllCoTyVarBndrUsing sym sco (Subst in_scope idenv tenv cenv) old_var old_kind_co
+substForAllCoBndrUsing :: CoOrDCo kco
+                       -> Bool  -- apply sym to binder?
+                       -> (Type -> Type)
+                       -> (kco -> kco)  -- transformation to kind co
+                       -> Subst -> TyCoVar -> kco
+                       -> (Subst, TyCoVar, kco)
+substForAllCoBndrUsing co_or_dco sym sty sco subst old_var
+  | isTyVar old_var = substForAllCoTyVarBndrUsing co_or_dco sym sty sco subst old_var
+  | otherwise       = substForAllCoCoVarBndrUsing co_or_dco sym sty sco subst old_var
+
+substForAllCoTyVarBndrUsing :: CoOrDCo kco
+                            -> Bool  -- apply sym to binder?
+                            -> (Type -> Type) -- transformation to types
+                            -> (kco -> kco)  -- transformation to kind co
+                            -> Subst -> TyVar -> kco
+                            -> (Subst, TyVar, kco)
+substForAllCoTyVarBndrUsing co_or_dco sym sty sco (Subst in_scope idenv tenv cenv) old_var old_kind_co
   = assert (isTyVar old_var )
     ( Subst (in_scope `extendInScopeSet` new_var) idenv new_env cenv
     , new_var, new_kind_co )
   where
     new_env | no_change && not sym = delVarEnv tenv old_var
             | sym       = extendVarEnv tenv old_var $
-                          TyVarTy new_var `CastTy` new_kind_co
+                          mk_cast (TyVarTy new_var) new_kind_co
             | otherwise = extendVarEnv tenv old_var (TyVarTy new_var)
 
-    no_kind_change = noFreeVarsOfCo old_kind_co
+    no_kind_change = case co_or_dco of
+      Co  -> noFreeVarsOfCo old_kind_co
+      DCo -> noFreeVarsOfDCo old_kind_co
+    mk_cast = case co_or_dco of
+      Co  -> CastTy
+      DCo -> pprPanic "substForAllCoTyVarBndrUsing DCo Sym"
+              (vcat [ text "kind_co:" <+> ppr old_kind_co ])
+
     no_change = no_kind_change && (new_var == old_var)
 
     new_kind_co | no_kind_change = old_kind_co
                 | otherwise      = sco old_kind_co
 
-    new_ki1 = coercionLKind new_kind_co
+    new_ki1 = case co_or_dco of
+      Co  -> coercionLKind new_kind_co
     -- We could do substitution to (tyVarKind old_var). We don't do so because
     -- we already substituted new_kind_co, which contains the kind information
     -- we want. We don't want to do substitution once more. Also, in most cases,
     -- new_kind_co is a Refl, in which case coercionKind is really fast.
+      DCo -> sty (tyVarKind old_var)
 
     new_var  = uniqAway in_scope (setTyVarKind old_var new_ki1)
 
-substForAllCoCoVarBndrUsing :: Bool  -- apply sym to binder?
-                            -> (Coercion -> Coercion)  -- transformation to kind co
-                            -> Subst -> CoVar -> KindCoercion
-                            -> (Subst, CoVar, KindCoercion)
-substForAllCoCoVarBndrUsing sym sco (Subst in_scope idenv tenv cenv)
+substForAllCoCoVarBndrUsing :: CoOrDCo kco
+                            -> Bool  -- apply sym to binder?
+                            -> (Type -> Type) -- transformation to types
+                            -> (kco -> kco)  -- transformation to kind co
+                            -> Subst -> CoVar -> kco
+                            -> (Subst, CoVar, kco)
+substForAllCoCoVarBndrUsing co_or_dco sym sty sco (Subst in_scope idenv tenv cenv)
                             old_var old_kind_co
-  = assert (isCoVar old_var )
+  = assert (isCoVar old_var)
     ( Subst (in_scope `extendInScopeSet` new_var) idenv tenv new_cenv
     , new_var, new_kind_co )
   where
     new_cenv | no_change && not sym = delVarEnv cenv old_var
              | otherwise = extendVarEnv cenv old_var (mkCoVarCo new_var)
 
-    no_kind_change = noFreeVarsOfCo old_kind_co
+    no_kind_change = case co_or_dco of
+      Co  -> noFreeVarsOfCo old_kind_co
+      DCo -> noFreeVarsOfDCo old_kind_co
     no_change = no_kind_change && (new_var == old_var)
 
     new_kind_co | no_kind_change = old_kind_co
                 | otherwise      = sco old_kind_co
 
-    Pair h1 h2 = coercionKind new_kind_co
+    Pair h1 h2 = case co_or_dco of
+      Co  -> coercionKind new_kind_co
+      DCo ->
+        let l_ty = sty (varType old_var)
+            r_ty = pprPanic "substForAllCoCoVarBndrUsing DCo Sym"
+                   (vcat [ text "kind_co:" <+> ppr old_kind_co])
+        in Pair l_ty r_ty
 
     new_var       = uniqAway in_scope $ mkCoVar (varName old_var) new_var_type
     new_var_type  | sym       = h2