[Git][ghc/ghc][wip/T24978] Better coercionKind

Simon Peyton Jones (@simonpj) gitlab at gitlab.haskell.org
Sun Jul 7 22:14:04 UTC 2024



Simon Peyton Jones pushed to branch wip/T24978 at Glasgow Haskell Compiler / GHC


Commits:
06af46d3 by Simon Peyton Jones at 2024-07-07T23:13:09+01:00
Better coercionKind

In particular

* Avoid #25066
* Do coercionLKind, coercionRKing by specialisation rather than duplication

- - - - -


4 changed files:

- compiler/GHC/Core/Coercion.hs
- compiler/GHC/Core/Coercion.hs-boot
- compiler/GHC/Core/DataCon.hs
- compiler/GHC/Core/TyCo/Subst.hs


Changes:

=====================================
compiler/GHC/Core/Coercion.hs
=====================================
@@ -24,7 +24,7 @@ module GHC.Core.Coercion (
 
         -- ** Functions over coercions
         coVarRType, coVarLType, coVarTypes,
-        coVarKind, coVarKindsTypesRole, coVarRole,
+        coVarKind, coVarTypesRole, coVarRole,
         coercionType, mkCoercionType,
         coercionKind, coercionLKind, coercionRKind,coercionKinds,
         coercionRole, coercionKindRole,
@@ -581,20 +581,18 @@ splitForAllCo_co_maybe _ = Nothing
 -- and some coercion kind stuff
 
 coVarLType, coVarRType :: HasDebugCallStack => CoVar -> Type
-coVarLType cv | (_, _, ty1, _, _) <- coVarKindsTypesRole cv = ty1
-coVarRType cv | (_, _, _, ty2, _) <- coVarKindsTypesRole cv = ty2
+coVarLType cv | (ty1, _, _) <- coVarTypesRole cv = ty1
+coVarRType cv | (_, ty2, _) <- coVarTypesRole cv = ty2
 
 coVarTypes :: HasDebugCallStack => CoVar -> Pair Type
-coVarTypes cv
-  | (_, _, ty1, ty2, _) <- coVarKindsTypesRole cv
-  = Pair ty1 ty2
-
-coVarKindsTypesRole :: HasDebugCallStack => CoVar -> (Kind,Kind,Type,Type,Role)
-coVarKindsTypesRole cv
- | Just (tc, [k1,k2,ty1,ty2]) <- splitTyConApp_maybe (varType cv)
- = (k1, k2, ty1, ty2, eqTyConRole tc)
+coVarTypes cv | (ty1, ty2, _) <- coVarTypesRole cv = Pair ty1 ty2
+
+coVarTypesRole :: HasDebugCallStack => CoVar -> (Type,Type,Role)
+coVarTypesRole cv
+ | Just (tc, [_,_,ty1,ty2]) <- splitTyConApp_maybe (varType cv)
+ = (ty1, ty2, eqTyConRole tc)
  | otherwise
- = pprPanic "coVarKindsTypesRole, non coercion variable"
+ = pprPanic "coVarTypesRole, non coercion variable"
             (ppr cv $$ ppr (varType cv))
 
 coVarKind :: CoVar -> Type
@@ -2091,7 +2089,7 @@ extendLiftingContextEx lc@(LC subst env) ((v,ty):rest)
     -- lift_s2 :: s2 ~r s2'
     -- kco     :: (s1 ~r s2) ~N (s1' ~r s2')
     assert (isCoVar v) $
-    let (_, _, s1, s2, r) = coVarKindsTypesRole v
+    let (s1, s2, r) = coVarTypesRole v
         lift_s1 = ty_co_subst lc r s1
         lift_s2 = ty_co_subst lc r s2
         kco     = mkTyConAppCo Nominal (equalityTyCon r)
@@ -2442,23 +2440,26 @@ coercionType co = case coercionKindRole co of
 coercionKind :: HasDebugCallStack => Coercion -> Pair Type
 coercionKind co = Pair (coercionLKind co) (coercionRKind co)
 
-coercionLKind :: HasDebugCallStack => Coercion -> Type
-coercionLKind co
+coercionLKind, coercionRKind :: HasDebugCallStack => Coercion -> Type
+coercionLKind co = coercion_lr_kind CLeft  co
+coercionRKind co = coercion_lr_kind CRight co
+
+coercion_lr_kind :: HasDebugCallStack => LeftOrRight -> Coercion -> Type
+{-# INLINE coercion_lr_kind #-}
+coercion_lr_kind which co
   = go co
   where
     go (Refl ty)                 = ty
     go (GRefl _ ty _)            = ty
     go (TyConAppCo _ tc cos)     = mkTyConApp tc (map go cos)
     go (AppCo co1 co2)           = mkAppTy (go co1) (go co2)
-    go (ForAllCo { fco_tcv = tv1, fco_visL = visL, fco_body = co1 })
-                                 = mkTyCoForAllTy tv1 visL (go co1)
     go (FunCo { fco_afl = af, fco_mult = mult, fco_arg = arg, fco_res = res})
        {- See Note [FunCo] -}    = FunTy { ft_af = af, ft_mult = go mult
                                          , ft_arg = go arg, ft_res = go res }
-    go (CoVarCo cv)              = coVarLType cv
-    go (HoleCo h)                = coVarLType (coHoleCoVar h)
+    go (CoVarCo cv)              = go_covar cv
+    go (HoleCo h)                = go_covar (coHoleCoVar h)
     go (UnivCo { uco_lty = ty1}) = ty1
-    go (SymCo co)                = coercionRKind co
+    go (SymCo co)                = pickLR which (coercionRKind co, coercionLKind co)
     go (TransCo co1 _)           = go co1
     go (LRCo lr co)              = pickLR lr (splitAppTy (go co))
     go (InstCo aco arg)          = go_app aco [go arg]
@@ -2467,77 +2468,66 @@ coercionLKind co
     go (SelCo d co)              = selectFromType d (go co)
     go (AxiomRuleCo ax cos)      = go_ax ax cos
 
+    go co@(ForAllCo { fco_tcv = tv1, fco_visL = visL, fco_visR = visR
+                    , fco_kind = k_co, fco_body = co1 })
+      = case which of
+          CLeft  -> mkTyCoForAllTy tv1 visL (go co1)
+          CRight | isGReflCo k_co  -- kind_co always has kind `Type`, thus `isGReflCo`
+                 -> mkTyCoForAllTy tv1 visR (go co1)
+                 | otherwise
+                 -> go_forall_right empty_subst co
+      where
+         empty_subst = mkEmptySubst (mkInScopeSet $ tyCoVarsOfCo co)
+
+    -------------
+    go_covar cv = pickLR which (coVarLType cv, coVarRType cv)
+
+    -------------
     go_app :: Coercion -> [Type] -> Type
     -- Collect up all the arguments and apply all at once
     -- See Note [Nested InstCos]
     go_app (InstCo co arg) args = go_app co (go arg:args)
     go_app co              args = piResultTys (go co) args
 
-    -- Tiresome to near-duplicate this, but coercionKind
-    -- is a very hot cod path
+    -------------
     go_ax (BuiltInFamRewrite  bif) cos = go_bif (bifrw_proves bif)  cos
     go_ax (BuiltInFamInteract bif) cos = go_bif (bifint_proves bif) cos
     go_ax (UnbranchedAxiom ax)     cos = go_branch ax (coAxiomSingleBranch ax) cos
     go_ax (BranchedAxiom ax i)     cos = go_branch ax (coAxiomNthBranch ax i)  cos
 
-    go_bif proves cos = pFst $ expectJust "coAxRuleKind" $ proves $
-                        map coercionKind cos
+    -------------
+    go_bif proves cos = case proves (map coercionKind cos) of
+                          Nothing -> pprPanic "coercionKind" (ppr cos)
+                          Just (Pair lty rty) -> pickLR which (lty, rty)
 
+    -------------
     go_branch :: CoAxiom br -> CoAxBranch -> [Coercion] -> Type
-    go_branch ax (CoAxBranch { cab_tvs = tvs, cab_cvs = cvs, cab_lhs = lhs_tys }) cos
+    go_branch ax (CoAxBranch { cab_tvs = tvs, cab_cvs = cvs
+                             , cab_lhs = lhs_tys, cab_rhs = rhs_ty }) cos
       = assert (cos `equalLength` tcvs) $
                   -- Invariant of AxiomRuleCo: cos should
                   -- exactly saturate the axiom branch
-        substTy (zipTCvSubst tcvs ltys) (mkTyConApp tc lhs_tys)
+        let (tys1, cotys1) = splitAtList tvs tys
+            cos1           = map stripCoercionTy cotys1
+        in
+-- You might think to use
+--        substTy (zipTCvSubst tcvs ltys) (mkTyConApp tc lhs_tys)
+-- but #25066 makes it much less efficient than the silly calls below
+        substTyWith tvs tys1       $
+        substTyWithCoVars cvs cos1 $
+        pickLR which (mkTyConApp tc lhs_tys, rhs_ty)
      where
        tc   = coAxiomTyCon ax
        tcvs | null cvs  = tvs  -- Very common case (currently always!)
             | otherwise = tvs ++ cvs
-       ltys = map go cos  -- `go` is `coercionLKind`
-
-coercionRKind :: HasDebugCallStack => Coercion -> Type
-coercionRKind co
-  = go co
-  where
-    go (Refl ty)                 = ty
-    go (GRefl _ ty MRefl)        = ty
-    go (GRefl _ ty (MCo co1))    = mkCastTy ty co1
-    go (TyConAppCo _ tc cos)     = mkTyConApp tc (map go cos)
-    go (AppCo co1 co2)           = mkAppTy (go co1) (go co2)
-    go (CoVarCo cv)              = coVarRType cv
-    go (HoleCo h)                = coVarRType (coHoleCoVar h)
-    go (FunCo { fco_afr = af, fco_mult = mult, fco_arg = arg, fco_res = res})
-       {- See Note [FunCo] -}    = FunTy { ft_af = af, ft_mult = go mult
-                                         , ft_arg = go arg, ft_res = go res }
-    go (UnivCo { uco_rty = ty2 })= ty2
-    go (SymCo co)                = coercionLKind co
-    go (TransCo _ co2)           = go co2
-    go (LRCo lr co)              = pickLR lr (splitAppTy (go co))
-    go (InstCo aco arg)          = go_app aco [go arg]
-    go (KindCo co)               = typeKind (go co)
-    go (SubCo co)                = go co
-    go (SelCo d co)              = selectFromType d (go co)
-    go (AxiomRuleCo ax cos)      = go_ax ax cos
+       tys = map go cos
 
-    go co@(ForAllCo { fco_tcv = tv1, fco_visR = visR
-                    , fco_kind = k_co, fco_body = co1 }) -- works for both tyvar and covar
-       | isGReflCo k_co           = mkTyCoForAllTy tv1 visR (go co1)
-         -- kind_co always has kind @Type@, thus @isGReflCo@
-       | otherwise                = go_forall empty_subst co
-       where
-         empty_subst = mkEmptySubst (mkInScopeSet $ tyCoVarsOfCo co)
-
-    go_app :: Coercion -> [Type] -> Type
-    -- Collect up all the arguments and apply all at once
-    -- See Note [Nested InstCos]
-    go_app (InstCo co arg) args = go_app co (go arg:args)
-    go_app co              args = piResultTys (go co) args
-
-    go_forall subst (ForAllCo { fco_tcv = tv1, fco_visR = visR
-                              , fco_kind = k_co, fco_body = co })
+    -------------
+    go_forall_right subst (ForAllCo { fco_tcv = tv1, fco_visR = visR
+                                    , fco_kind = k_co, fco_body = co })
       -- See Note [Nested ForAllCos]
       | isTyVar tv1
-      = mkForAllTy (Bndr tv2 visR) (go_forall subst' co)
+      = mkForAllTy (Bndr tv2 visR) (go_forall_right subst' co)
       where
         k2  = coercionRKind k_co
         tv2 = setTyVarKind tv1 (substTy subst k2)
@@ -2546,10 +2536,10 @@ coercionRKind co
                | otherwise      = extendTvSubst (extendSubstInScope subst tv2) tv1 $
                                   TyVarTy tv2 `mkCastTy` mkSymCo k_co
 
-    go_forall subst (ForAllCo { fco_tcv = cv1, fco_visR = visR
-                              , fco_kind = k_co, fco_body = co })
+    go_forall_right subst (ForAllCo { fco_tcv = cv1, fco_visR = visR
+                                    , fco_kind = k_co, fco_body = co })
       | isCoVar cv1
-      = mkTyCoForAllTy cv2 visR (go_forall subst' co)
+      = mkTyCoForAllTy cv2 visR (go_forall_right subst' co)
       where
         k2    = coercionRKind k_co
         r     = coVarRole cv1
@@ -2572,30 +2562,10 @@ coercionRKind co
                 | otherwise     = extendCvSubst (extendSubstInScope subst cv2)
                                                 cv1 n_subst
 
-    go_forall subst other_co
+    go_forall_right subst other_co
       -- when other_co is not a ForAllCo
       = substTy subst (go other_co)
 
-    -- Tiresome to near-duplicate this, but coercionKind
-    -- is a very hot cod path
-    go_ax (BuiltInFamRewrite  bif) cos = go_bif (bifrw_proves bif)  cos
-    go_ax (BuiltInFamInteract bif) cos = go_bif (bifint_proves bif) cos
-    go_ax (UnbranchedAxiom ax)     cos = go_branch (coAxiomSingleBranch ax) cos
-    go_ax (BranchedAxiom ax i)     cos = go_branch (coAxiomNthBranch ax i)  cos
-
-    go_bif proves cos = pSnd $ expectJust "coAxRuleKind" $ proves $
-                        map coercionKind cos
-
-    go_branch :: CoAxBranch -> [Coercion] -> Type
-    go_branch (CoAxBranch { cab_tvs = tvs, cab_cvs = cvs, cab_rhs = rhs_ty }) cos
-      = assert (cos `equalLength` tcvs) $
-                  -- Invariant of AxiomRuleCo: cos should
-                  -- exactly saturate the axiom branch
-        substTy (zipTCvSubst tcvs rtys) rhs_ty
-     where
-       tcvs | null cvs  = tvs  -- Very common case (currently always!)
-            | otherwise = tvs ++ cvs
-       rtys = map go cos  -- `go` is `coercionRKind`
 {-
 
 Note [Nested ForAllCos]


=====================================
compiler/GHC/Core/Coercion.hs-boot
=====================================
@@ -41,7 +41,7 @@ isGReflCo :: Coercion -> Bool
 isReflCo :: Coercion -> Bool
 isReflexiveCo :: Coercion -> Bool
 decomposePiCos :: HasDebugCallStack => Coercion -> Pair Type -> [Type] -> ([Coercion], Coercion)
-coVarKindsTypesRole :: HasDebugCallStack => CoVar -> (Kind, Kind, Type, Type, Role)
+coVarTypesRole :: HasDebugCallStack => CoVar -> (Type, Type, Role)
 coVarRole :: CoVar -> Role
 
 mkCoercionType :: Role -> Type -> Type -> Type


=====================================
compiler/GHC/Core/DataCon.hs
=====================================
@@ -1338,7 +1338,7 @@ dataConKindEqSpec (MkData {dcExTyCoVars = ex_tcvs})
   = [ EqSpec tv ty
     | cv <- ex_tcvs
     , isCoVar cv
-    , let (_, _, ty1, ty, _) = coVarKindsTypesRole cv
+    , let (ty1, ty, _) = coVarTypesRole cv
           tv = getTyVar ty1
     ]
 


=====================================
compiler/GHC/Core/TyCo/Subst.hs
=====================================
@@ -60,7 +60,7 @@ import {-# SOURCE #-} GHC.Core.Coercion
    , mkAxiomRuleCo, mkAppCo, mkGReflCo
    , mkInstCo, mkLRCo, mkTyConAppCo
    , mkCoercionType
-   , coercionKind, coercionLKind, coVarKindsTypesRole )
+   , coercionKind, coercionLKind, coVarTypesRole )
 import {-# SOURCE #-} GHC.Core.TyCo.Ppr ( pprTyVar )
 import {-# SOURCE #-} GHC.Core.Ppr ( ) -- instance Outputable CoreExpr
 import {-# SOURCE #-} GHC.Core ( CoreExpr )
@@ -1086,7 +1086,7 @@ substCoVarBndrUsing subst_fn subst@(Subst in_scope idenv tenv cenv) old_var
     new_var = uniqAway in_scope subst_old_var
     subst_old_var = mkCoVar (varName old_var) new_var_type
 
-    (_, _, t1, t2, role) = coVarKindsTypesRole old_var
+    (t1, t2, role) = coVarTypesRole old_var
     t1' = subst_fn subst t1
     t2' = subst_fn subst t2
     new_var_type = mkCoercionType role t1' t2'



View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/06af46d317de4d016ece6a699b9f5a042863f41a

-- 
This project does not include diff previews in email notifications.
View it on GitLab: https://gitlab.haskell.org/ghc/ghc/-/commit/06af46d317de4d016ece6a699b9f5a042863f41a
You're receiving this email because of your account on gitlab.haskell.org.


-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.haskell.org/pipermail/ghc-commits/attachments/20240707/c50452cb/attachment-0001.html>


More information about the ghc-commits mailing list