Caching coercion roles in NthCo (#11735)

13 files changed, 64 insertions, 43 deletions

diff --git a/compiler/coreSyn/CoreFVs.hs b/compiler/coreSyn/CoreFVs.hs
index 7fcff90c71..7e7727164c 100644
--- a/compiler/coreSyn/CoreFVs.hs
+++ b/compiler/coreSyn/CoreFVs.hs
@@ -379,7 +379,7 @@ orphNamesOfCo (AxiomInstCo con _ cos) = orphNamesOfCoCon con `unionNameSet` orph
 orphNamesOfCo (UnivCo p _ t1 t2)    = orphNamesOfProv p `unionNameSet` orphNamesOfType t1 `unionNameSet` orphNamesOfType t2
 orphNamesOfCo (SymCo co)            = orphNamesOfCo co
 orphNamesOfCo (TransCo co1 co2)     = orphNamesOfCo co1 `unionNameSet` orphNamesOfCo co2
-orphNamesOfCo (NthCo _ co)          = orphNamesOfCo co
+orphNamesOfCo (NthCo _ _ co)        = orphNamesOfCo co
 orphNamesOfCo (LRCo  _ co)          = orphNamesOfCo co
 orphNamesOfCo (InstCo co arg)       = orphNamesOfCo co `unionNameSet` orphNamesOfCo arg
 orphNamesOfCo (CoherenceCo co1 co2) = orphNamesOfCo co1 `unionNameSet` orphNamesOfCo co2
diff --git a/compiler/coreSyn/CoreLint.hs b/compiler/coreSyn/CoreLint.hs
index e866f0d3ec..ce3054aaff 100644
--- a/compiler/coreSyn/CoreLint.hs
+++ b/compiler/coreSyn/CoreLint.hs
@@ -1731,7 +1731,7 @@ lintCoercion co@(TransCo co1 co2)
        ; lintRole co r1 r2
        ; return (k1a, k2b, ty1a, ty2b, r1) }
 
-lintCoercion the_co@(NthCo n co)
+lintCoercion the_co@(NthCo _ n co)
   = do { (_, _, s, t, r) <- lintCoercion co
        ; case (splitForAllTy_maybe s, splitForAllTy_maybe t) of
          { (Just (tv_s, _ty_s), Just (tv_t, _ty_t))
diff --git a/compiler/iface/TcIface.hs b/compiler/iface/TcIface.hs
index 70438f6337..6ac006920d 100644
--- a/compiler/iface/TcIface.hs
+++ b/compiler/iface/TcIface.hs
@@ -1353,7 +1353,7 @@ tcIfaceCo = go
                                             <*> go c2
     go (IfaceInstCo c1 t2)       = InstCo   <$> go c1
                                             <*> go t2
-    go (IfaceNthCo d c)          = NthCo d  <$> go c
+    go (IfaceNthCo d c)          = mkNthCo d  <$> go c
     go (IfaceLRCo lr c)          = LRCo lr  <$> go c
     go (IfaceCoherenceCo c1 c2)  = CoherenceCo <$> go c1
                                                <*> go c2
diff --git a/compiler/iface/ToIface.hs b/compiler/iface/ToIface.hs
index da24d38f9e..b22910bf2b 100644
--- a/compiler/iface/ToIface.hs
+++ b/compiler/iface/ToIface.hs
@@ -224,7 +224,7 @@ toIfaceCoercionX fr co
     go (AppCo co1 co2)      = IfaceAppCo  (go co1) (go co2)
     go (SymCo co)           = IfaceSymCo (go co)
     go (TransCo co1 co2)    = IfaceTransCo (go co1) (go co2)
-    go (NthCo d co)         = IfaceNthCo d (go co)
+    go (NthCo _ d co)       = IfaceNthCo d (go co)
     go (LRCo lr co)         = IfaceLRCo lr (go co)
     go (InstCo co arg)      = IfaceInstCo (go co) (go arg)
     go (CoherenceCo c1 c2)  = IfaceCoherenceCo (go c1) (go c2)
diff --git a/compiler/typecheck/TcTyDecls.hs b/compiler/typecheck/TcTyDecls.hs
index 548f058811..2f6387ac0e 100644
--- a/compiler/typecheck/TcTyDecls.hs
+++ b/compiler/typecheck/TcTyDecls.hs
@@ -122,7 +122,7 @@ synonymTyConsOfType ty
      go_co (UnivCo p _ ty ty')    = go_prov p `plusNameEnv` go ty `plusNameEnv` go ty'
      go_co (SymCo co)             = go_co co
      go_co (TransCo co co')       = go_co co `plusNameEnv` go_co co'
-     go_co (NthCo _ co)           = go_co co
+     go_co (NthCo _ _ co)         = go_co co
      go_co (LRCo _ co)            = go_co co
      go_co (InstCo co co')        = go_co co `plusNameEnv` go_co co'
      go_co (CoherenceCo co co')   = go_co co `plusNameEnv` go_co co'
diff --git a/compiler/typecheck/TcType.hs b/compiler/typecheck/TcType.hs
index de37aa862d..4f77312242 100644
--- a/compiler/typecheck/TcType.hs
+++ b/compiler/typecheck/TcType.hs
@@ -908,7 +908,7 @@ exactTyCoVarsOfType ty
     goCo (UnivCo p _ t1 t2)  = goProv p `unionVarSet` go t1 `unionVarSet` go t2
     goCo (SymCo co)          = goCo co
     goCo (TransCo co1 co2)   = goCo co1 `unionVarSet` goCo co2
-    goCo (NthCo _ co)        = goCo co
+    goCo (NthCo _ _ co)      = goCo co
     goCo (LRCo _ co)         = goCo co
     goCo (InstCo co arg)     = goCo co `unionVarSet` goCo arg
     goCo (CoherenceCo c1 c2) = goCo c1 `unionVarSet` goCo c2
diff --git a/compiler/typecheck/TcUnify.hs b/compiler/typecheck/TcUnify.hs
index 2c374285fc..155bea6369 100644
--- a/compiler/typecheck/TcUnify.hs
+++ b/compiler/typecheck/TcUnify.hs
@@ -2135,7 +2135,7 @@ occCheckExpand tv ty
     go_co env (TransCo co1 co2)         = do { co1' <- go_co env co1
                                              ; co2' <- go_co env co2
                                              ; return (mkTransCo co1' co2') }
-    go_co env (NthCo n co)              = do { co' <- go_co env co
+    go_co env (NthCo r n co)            = do { co' <- go_co env co -- CHECK
                                              ; return (mkNthCo n co') }
     go_co env (LRCo lr co)              = do { co' <- go_co env co
                                              ; return (mkLRCo lr co') }
diff --git a/compiler/typecheck/TcValidity.hs b/compiler/typecheck/TcValidity.hs
index 8c01460f2e..f30d893232 100644
--- a/compiler/typecheck/TcValidity.hs
+++ b/compiler/typecheck/TcValidity.hs
@@ -1964,7 +1964,7 @@ fvCo (AxiomInstCo _ _ args) = concatMap fvCo args
 fvCo (UnivCo p _ t1 t2)     = fvProv p ++ fvType t1 ++ fvType t2
 fvCo (SymCo co)             = fvCo co
 fvCo (TransCo co1 co2)      = fvCo co1 ++ fvCo co2
-fvCo (NthCo _ co)           = fvCo co
+fvCo (NthCo _ _ co)         = fvCo co
 fvCo (LRCo _ co)            = fvCo co
 fvCo (InstCo co arg)        = fvCo co ++ fvCo arg
 fvCo (CoherenceCo co1 co2)  = fvCo co1 ++ fvCo co2
diff --git a/compiler/types/Coercion.hs b/compiler/types/Coercion.hs
index 770c94f4d6..c600072e78 100644
--- a/compiler/types/Coercion.hs
+++ b/compiler/types/Coercion.hs
@@ -846,7 +846,25 @@ mkNthCo n co@(FunCo _ arg res)
 
 mkNthCo n (TyConAppCo _ _ arg_cos) = arg_cos `getNth` n
 
-mkNthCo n co = NthCo n co
+mkNthCo n co =
+  NthCo r n co
+  where
+    r = nthCoercionRole co
+
+    nthCoercionRole co
+      | Just (tv1, _) <- splitForAllTy_maybe ty1
+      = ASSERT( n == 0 )
+        Nominal
+
+      | otherwise
+      = let (tc1,  args1) = splitTyConApp ty1
+            (_tc2, args2) = splitTyConApp ty2
+        in
+        ASSERT2( tc1 == _tc2, ppr n $$ ppr tc1 $$ ppr _tc2 )
+        (nthRole r tc1 n)
+
+      where
+        (Pair ty1 ty2, r) = coercionKindRole co
 
 mkLRCo :: LeftOrRight -> Coercion -> Coercion
 mkLRCo lr (Refl eq ty) = Refl eq (pickLR lr (splitAppTy ty))
@@ -994,8 +1012,8 @@ setNominalRole_maybe (AppCo co1 co2)
   = AppCo <$> setNominalRole_maybe co1 <*> pure co2
 setNominalRole_maybe (ForAllCo tv kind_co co)
   = ForAllCo tv kind_co <$> setNominalRole_maybe co
-setNominalRole_maybe (NthCo n co)
-  = NthCo n <$> setNominalRole_maybe co
+setNominalRole_maybe (NthCo r n co)
+  = NthCo r n <$> setNominalRole_maybe co
 setNominalRole_maybe (InstCo co arg)
   = InstCo <$> setNominalRole_maybe co <*> pure arg
 setNominalRole_maybe (CoherenceCo co1 co2)
@@ -1104,7 +1122,7 @@ promoteCoercion co = case co of
     TransCo co1 co2
       -> mkTransCo (promoteCoercion co1) (promoteCoercion co2)
 
-    NthCo n co1
+    NthCo _ n co1
       | Just (_, args) <- splitTyConAppCo_maybe co1
       , args `lengthExceeds` n
       -> promoteCoercion (args !! n)
@@ -1642,7 +1660,7 @@ seqCo (UnivCo p r t1 t2)
   = seqProv p `seq` r `seq` seqType t1 `seq` seqType t2
 seqCo (SymCo co)                = seqCo co
 seqCo (TransCo co1 co2)         = seqCo co1 `seq` seqCo co2
-seqCo (NthCo n co)              = n `seq` seqCo co
+seqCo (NthCo r n co)            = n `seq` seqCo co
 seqCo (LRCo lr co)              = lr `seq` seqCo co
 seqCo (InstCo co arg)           = seqCo co `seq` seqCo arg
 seqCo (CoherenceCo co1 co2)     = seqCo co1 `seq` seqCo co2
@@ -1731,7 +1749,7 @@ coercionKind co = go co
     go (UnivCo _ _ ty1 ty2)   = Pair ty1 ty2
     go (SymCo co)             = swap $ go co
     go (TransCo co1 co2)      = Pair (pFst $ go co1) (pSnd $ go co2)
-    go g@(NthCo d co)
+    go g@(NthCo _ d co)
       | Just argss <- traverse tyConAppArgs_maybe tys
       = ASSERT( and $ (`lengthExceeds` d) <$> argss )
         (`getNth` d) <$> argss
@@ -1789,20 +1807,7 @@ coercionRole = go
     go (UnivCo _ r _ _)  = r
     go (SymCo co) = go co
     go (TransCo co1 co2) = go co1
-    go (NthCo d co)
-      | Just (tv1, _) <- splitForAllTy_maybe ty1
-      = ASSERT( d == 0 )
-        Nominal
-
-      | otherwise
-      = let (tc1,  args1) = splitTyConApp ty1
-            (_tc2, args2) = splitTyConApp ty2
-        in
-        ASSERT2( tc1 == _tc2, ppr d $$ ppr tc1 $$ ppr _tc2 )
-        (nthRole r tc1 d)
-
-      where
-        (Pair ty1 ty2, r) = coercionKindRole co
+    go (NthCo r d co) = r
     go (LRCo {}) = Nominal
     go (InstCo co arg) = go_app co
     go (CoherenceCo co1 _) = go co1
diff --git a/compiler/types/FamInstEnv.hs b/compiler/types/FamInstEnv.hs
index 131abda47b..1b9a6cb6a9 100644
--- a/compiler/types/FamInstEnv.hs
+++ b/compiler/types/FamInstEnv.hs
@@ -1634,7 +1634,7 @@ allTyVarsInTy = go
     go_co (UnivCo p _ t1 t2)    = go_prov p `unionVarSet` go t1 `unionVarSet` go t2
     go_co (SymCo co)            = go_co co
     go_co (TransCo c1 c2)       = go_co c1 `unionVarSet` go_co c2
-    go_co (NthCo _ co)          = go_co co
+    go_co (NthCo _ _ co)        = go_co co
     go_co (LRCo _ co)           = go_co co
     go_co (InstCo co arg)       = go_co co `unionVarSet` go_co arg
     go_co (CoherenceCo c1 c2)   = go_co c1 `unionVarSet` go_co c2
diff --git a/compiler/types/OptCoercion.hs b/compiler/types/OptCoercion.hs
index 24dc8a4344..5a020a4d2c 100644
--- a/compiler/types/OptCoercion.hs
+++ b/compiler/types/OptCoercion.hs
@@ -442,7 +442,7 @@ opt_univ env sym prov role oty1 oty2
 opt_nth_co :: LiftingContext -> SymFlag -> ReprFlag -> Role -> Coercion -> NormalCo
 opt_nth_co env sym rep r = go []
   where
-    go ns (NthCo n co) = go (n:ns) co
+    go ns (NthCo _ n co) = go (n:ns) co
       -- previous versions checked if the tycon is decomposable. This
       -- is redundant, because a non-decomposable tycon under an NthCo
       -- is entirely bogus. See docs/core-spec/core-spec.pdf.
@@ -529,7 +529,7 @@ opt_trans2 _ co1 co2
 opt_trans_rule :: InScopeSet -> NormalNonIdCo -> NormalNonIdCo -> Maybe NormalCo
 
 -- Push transitivity through matching destructors
-opt_trans_rule is in_co1@(NthCo d1 co1) in_co2@(NthCo d2 co2)
+opt_trans_rule is in_co1@(NthCo _ d1 co1) in_co2@(NthCo _ d2 co2)
   | d1 == d2
   , co1 `compatible_co` co2
   = fireTransRule "PushNth" in_co1 in_co2 $
diff --git a/compiler/types/TyCoRep.hs b/compiler/types/TyCoRep.hs
index abccc9d651..830c07c6e6 100644
--- a/compiler/types/TyCoRep.hs
+++ b/compiler/types/TyCoRep.hs
@@ -824,10 +824,11 @@ data Coercion
     -- of the CoAxiomRule (i.e., the rule is fully saturated).
   | AxiomRuleCo CoAxiomRule [Coercion]
 
-  | NthCo  Int         Coercion     -- Zero-indexed; decomposes (T t0 ... tn)
+  | NthCo  Role Int Coercion     -- Zero-indexed; decomposes (T t0 ... tn)
     -- :: _ -> e -> ?? (inverse of TyConAppCo, see Note [TyConAppCo roles])
     -- Using NthCo on a ForAllCo gives an N coercion always
     -- See Note [NthCo and newtypes]
+    -- See Note [NthCo Cached Roles]
 
   | LRCo   LeftOrRight CoercionN     -- Decomposes (t_left t_right)
     -- :: _ -> N -> N
@@ -1137,7 +1138,7 @@ We can then build
 for any `a` and `b`. Because of the role annotation on N, if we use
 NthCo, we'll get out a representational coercion. That is:
 
-  NthCo 0 co :: forall a b. a ~R b
+  NthCo r 0 co :: forall a b. a ~R b
 
 Yikes! Clearly, this is terrible. The solution is simple: forbid
 NthCo to be used on newtypes if the internal coercion is representational.
@@ -1146,6 +1147,17 @@ This is not just some corner case discovered by a segfault somewhere;
 it was discovered in the proof of soundness of roles and described
 in the "Safe Coercions" paper (ICFP '14).
 
+Note [NthCo Cached Roles]
+~~~~~~~~~~~~~~~~~~~~~~~~~
+Looking at coercionKind and coercionRole, we see that in order to find the Role
+for an NthCo, we have to know the kind; unfortunately, calculating the coercion
+kind for a nested NthCo is quadratic in the nesting depth, so we would like to
+avoid doing that as much as possible. To this end, we calculate the Role at
+construction time, and store it in the NthCo itself, thus removing the need to
+do the calculation on the fly.
+
+See Trac #11735 for the discussion that led to this.
+
 Note [InstCo roles]
 ~~~~~~~~~~~~~~~~~~~
 Here is (essentially) the typing rule for InstCo:
@@ -1491,7 +1503,7 @@ tyCoFVsOfCo (UnivCo p _ t1 t2) fv_cand in_scope acc
                      `unionFV` tyCoFVsOfType t2) fv_cand in_scope acc
 tyCoFVsOfCo (SymCo co)          fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
 tyCoFVsOfCo (TransCo co1 co2)   fv_cand in_scope acc = (tyCoFVsOfCo co1 `unionFV` tyCoFVsOfCo co2) fv_cand in_scope acc
-tyCoFVsOfCo (NthCo _ co)        fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
+tyCoFVsOfCo (NthCo _ _ co)      fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
 tyCoFVsOfCo (LRCo _ co)         fv_cand in_scope acc = tyCoFVsOfCo co fv_cand in_scope acc
 tyCoFVsOfCo (InstCo co arg)     fv_cand in_scope acc = (tyCoFVsOfCo co `unionFV` tyCoFVsOfCo arg) fv_cand in_scope acc
 tyCoFVsOfCo (CoherenceCo c1 c2) fv_cand in_scope acc = (tyCoFVsOfCo c1 `unionFV` tyCoFVsOfCo c2) fv_cand in_scope acc
@@ -1554,7 +1566,7 @@ coVarsOfCo (AxiomInstCo _ _ as) = coVarsOfCos as
 coVarsOfCo (UnivCo p _ t1 t2)   = coVarsOfProv p `unionVarSet` coVarsOfTypes [t1, t2]
 coVarsOfCo (SymCo co)           = coVarsOfCo co
 coVarsOfCo (TransCo co1 co2)    = coVarsOfCo co1 `unionVarSet` coVarsOfCo co2
-coVarsOfCo (NthCo _ co)         = coVarsOfCo co
+coVarsOfCo (NthCo _ _ co)       = coVarsOfCo co
 coVarsOfCo (LRCo _ co)          = coVarsOfCo co
 coVarsOfCo (InstCo co arg)      = coVarsOfCo co `unionVarSet` coVarsOfCo arg
 coVarsOfCo (CoherenceCo c1 c2)  = coVarsOfCos [c1, c2]
@@ -1661,7 +1673,7 @@ noFreeVarsOfCo (UnivCo p _ t1 t2)     = noFreeVarsOfProv p &&
                                         noFreeVarsOfType t2
 noFreeVarsOfCo (SymCo co)             = noFreeVarsOfCo co
 noFreeVarsOfCo (TransCo co1 co2)      = noFreeVarsOfCo co1 && noFreeVarsOfCo co2
-noFreeVarsOfCo (NthCo _ co)           = noFreeVarsOfCo co
+noFreeVarsOfCo (NthCo _ _ co)         = noFreeVarsOfCo co
 noFreeVarsOfCo (LRCo _ co)            = noFreeVarsOfCo co
 noFreeVarsOfCo (InstCo co1 co2)       = noFreeVarsOfCo co1 && noFreeVarsOfCo co2
 noFreeVarsOfCo (CoherenceCo co1 co2)  = noFreeVarsOfCo co1 && noFreeVarsOfCo co2
@@ -2330,7 +2342,11 @@ subst_co subst co
                                 (go_ty t1)) $! (go_ty t2)
     go (SymCo co)            = mkSymCo $! (go co)
     go (TransCo co1 co2)     = (mkTransCo $! (go co1)) $! (go co2)
-    go (NthCo d co)          = mkNthCo d $! (go co)
+    go (NthCo r d co)        = mkNthCo d $! (go co)
+    -- Q: Would it be reasonable to take the Role r and pass it to mkNthCoRole
+    -- here, instead of relying on mkNthCo to calculate the correct Role for 
+    -- us?
+
     go (LRCo lr co)          = mkLRCo lr $! (go co)
     go (InstCo co arg)       = (mkInstCo $! (go co)) $! go arg
     go (CoherenceCo co1 co2) = (mkCoherenceCo $! (go co1)) $! (go co2)
@@ -2930,7 +2946,7 @@ tidyCo env@(_, subst) co
                                 tidyType env t1) $! tidyType env t2
     go (SymCo co)            = SymCo $! go co
     go (TransCo co1 co2)     = (TransCo $! go co1) $! go co2
-    go (NthCo d co)          = NthCo d $! go co
+    go (NthCo r d co)        = NthCo r d $! go co
     go (LRCo lr co)          = LRCo lr $! go co
     go (InstCo co ty)        = (InstCo $! go co) $! go ty
     go (CoherenceCo co1 co2) = (CoherenceCo $! go co1) $! go co2
@@ -2989,7 +3005,7 @@ coercionSize (AxiomInstCo _ _ args) = 1 + sum (map coercionSize args)
 coercionSize (UnivCo p _ t1 t2)  = 1 + provSize p + typeSize t1 + typeSize t2
 coercionSize (SymCo co)          = 1 + coercionSize co
 coercionSize (TransCo co1 co2)   = 1 + coercionSize co1 + coercionSize co2
-coercionSize (NthCo _ co)        = 1 + coercionSize co
+coercionSize (NthCo _ _ co)      = 1 + coercionSize co
 coercionSize (LRCo  _ co)        = 1 + coercionSize co
 coercionSize (InstCo co arg)     = 1 + coercionSize co + coercionSize arg
 coercionSize (CoherenceCo c1 c2) = 1 + coercionSize c1 + coercionSize c2
diff --git a/compiler/types/Type.hs b/compiler/types/Type.hs
index 3f893dbcb2..ef80919aed 100644
--- a/compiler/types/Type.hs
+++ b/compiler/types/Type.hs
@@ -423,7 +423,7 @@ expandTypeSynonyms ty
       = mkSymCo (go_co subst co)
     go_co subst (TransCo co1 co2)
       = mkTransCo (go_co subst co1) (go_co subst co2)
-    go_co subst (NthCo n co)
+    go_co subst (NthCo _ n co)
       = mkNthCo n (go_co subst co)
     go_co subst (LRCo lr co)
       = mkLRCo lr (go_co subst co)
@@ -560,7 +560,7 @@ mapCoercion mapper@(TyCoMapper { tcm_smart = smart, tcm_covar = covar
     go (SymCo co) = mksymco <$> go co
     go (TransCo c1 c2) = mktransco <$> go c1 <*> go c2
     go (AxiomRuleCo r cos) = AxiomRuleCo r <$> mapM go cos
-    go (NthCo i co)        = mknthco i <$> go co
+    go (NthCo r i co)      = mknthco r i <$> go co
     go (LRCo lr co)        = mklrco lr <$> go co
     go (InstCo co arg)     = mkinstco <$> go co <*> go arg
     go (CoherenceCo c1 c2) = mkcoherenceco <$> go c1 <*> go c2
@@ -577,7 +577,7 @@ mapCoercion mapper@(TyCoMapper { tcm_smart = smart, tcm_covar = covar
       , mkkindco, mksubco, mkforallco)
       | smart
       = ( mkTyConAppCo, mkAppCo, mkAxiomInstCo, mkUnivCo
-        , mkSymCo, mkTransCo, mkNthCo, mkLRCo, mkInstCo, mkCoherenceCo
+        , mkSymCo, mkTransCo, const mkNthCo, mkLRCo, mkInstCo, mkCoherenceCo
         , mkKindCo, mkSubCo, mkForAllCo )
       | otherwise
       = ( TyConAppCo, AppCo, AxiomInstCo, UnivCo
@@ -2400,7 +2400,7 @@ tyConsOfType ty
      go_co (HoleCo {})             = emptyUniqSet
      go_co (SymCo co)              = go_co co
      go_co (TransCo co1 co2)       = go_co co1 `unionUniqSets` go_co co2
-     go_co (NthCo _ co)            = go_co co
+     go_co (NthCo _ _ co)          = go_co co
      go_co (LRCo _ co)             = go_co co
      go_co (InstCo co arg)         = go_co co `unionUniqSets` go_co arg
      go_co (CoherenceCo co1 co2)   = go_co co1 `unionUniqSets` go_co co2