A bit closer, but still stuckwip/T17323

I have spent far too long trying to implement our `ZonkCo` idea, but
so far I have failed.  The branch is `wip/T17523`.

Here are the things I'm stuck on:

**The occurs check**

In `TcCanonical.canEqTyVar2` we use `metaTyVarUpdateOK` to do an
occurs check. The latter explicitly recurses into the kind of any
variables, lest we are trying to unify
```
alpha ~ ((beta :: alpha -> *) (gamma :: alpha)
 ```
where alpha appears in the kind of beta and gamma. That is obviously a
bogus constraint, but we really, really don't want to unify alpha else
we'll build an infinite type and zonking will loop.
`metaTyVarUpdateOK` will spot this.

But suppose it was `alpha ~ ((beta :: delta -> *) (gamma :: delta)``
where we have already unified `delta := alpha`.  Now
metaTyVarUpdateOK` will not spot it!

ZonkCo aside, this looks like a bad bug.

It is hard to trigger.
* We could not have `alpha ~ (beta :: (..alpha...))` because the
  equality is homo-kinded and (I assume, somehow) a variable can't
  occur in its own kind.
* We could not have `alpha ~ T (beta ::(...alpha...))` because, since
  T's kind is closed, alpha would have to show up, visibly, in T's
  kind arguments, and they *are* flattened.

So I think it may just be `AppTy`.  Suppose we imagine a type
constructor `APP :: forall k1 k2. (k1->k2) -> k1 -> k2`.  Now I wonder
if the flattener might regard `AppTy t1 t2` as rather like `TyConApp
APP [k1, k2, t1, t2]`, where `t1:k1` and `t2:k2`.  If we really did
that, all would be well; indeed `metaTyVarUpdateOK` woudl not even
need to look in the kinds of type variables.

**zonk_eq_types**

`TcCanonical.zonk_eq_types` does a partial zonk of its argument type
as it goes; if the types are unequal, it returns the partially-zonked
types.  But partially zonking a type is going to invalidate the WKTI.
We must fully zonk or not at all.  Yikes.  I suppose we can throw away
our work in the un-equal case.  Or perhaps add ZonkCos internally in
teh returned types... Tricky.

I also noticed that in this code
```
can_eq_nc' False rdr_env envs ev eq_rel _ ps_ty1 _ ps_ty2
  = do { (xi1, co1) <- flatten FM_FlattenAll ev ps_ty1
       ; (xi2, co2) <- flatten FM_FlattenAll ev ps_ty2
       ; new_ev <- rewriteEqEvidence ev NotSwapped xi1 xi2 co1 co2
       ; can_eq_nc' True rdr_env envs new_ev eq_rel xi1 xi1 xi2 xi2 }
```
we discard `t1` and `t2`.  Yet they are the carefully-preserved
outputs of `zonk_eq_types`, so we end up repeating the work anyway!!
Maybe we should flatten `t1` and `t2`; but that might lose useful
type-synonym information.

**Unflattening**

I think it's very important to have the invariant that, when unifying
`alpha := ty`, alpha's kind is equal to ty's.  Equal, not just that
they zonk to the same thing.  Otherwise the flattener will no longer
be kind-homogeneous.

But the unflattener does unifications that don't obey this: it does a
`zonkTcType` before unifying.  Why?  That pesky occurs check again!

**The type of evidence Ids**

See `TcSMonad.add_item`:
```
      -- The evidence Id in a Given equality in InertEqs must match
      -- the predicate (which is not true of all CTyEqCans), else when
      -- we fetch it out in TcFlatten.flatten_tyvar2, ctEvCoercion won't have the
      -- correct type, and we'll mess up the WKTI
      -- See Note [Ct/evidence invariant]
```
I think this does fix the problem, but it's pretty tricky.

5 files changed, 87 insertions, 62 deletions

diff --git a/compiler/typecheck/TcCanonical.hs b/compiler/typecheck/TcCanonical.hs
index 5489e4d2cc..5b7f50ccd6 100644
--- a/compiler/typecheck/TcCanonical.hs
+++ b/compiler/typecheck/TcCanonical.hs
@@ -1362,10 +1362,22 @@ canEqCast flat ev eq_rel swapped ty1 co1 ty2 ps_ty2
   = do { traceTcS "Decomposing cast" (vcat [ ppr ev
                                            , ppr ty1 <+> text "|>" <+> ppr co1
                                            , ppr ps_ty2 ])
-       ; new_ev <- rewriteEqEvidence ev swapped ty1 ps_ty2
-                                     (mkTcGReflRightCo role ty1 co1)
-                                     (mkTcReflCo role ps_ty2)
-       ; can_eq_nc flat new_ev eq_rel ty1 ty1 ty2 ps_ty2 }
+       ; let lhs_co = mkTcGReflRightCo role ty1 co1
+             rhs_co = mkTcReflCo role ps_ty2
+
+       -- Under NotSwapped we have
+       --   ev     :: (ty1 |> co1) ~ ty2
+       --   co1    :: kind(ty1) ~ k2
+       --   lhs_co :: ty1 ~ (ty1 |> co1)
+       --   rhs_co :: ty2 ~ ty2
+       --   new_ev :: ty1 ~ ty2
+       ; case swapped of   -- Swap back because we want to hit the
+                           -- Note [Rewriting with ZonkCo] case of rewriteEqEvidence
+          NotSwapped -> do { new_ev <- rewriteEqEvidence ev NotSwapped ty1 ps_ty2 lhs_co rhs_co
+                           ; can_eq_nc flat new_ev eq_rel ty1 ty1 ty2 ps_ty2 }
+          IsSwapped  -> do { new_ev <- rewriteEqEvidence ev NotSwapped ps_ty2 ty1 rhs_co lhs_co
+                           ; can_eq_nc flat new_ev eq_rel ty2 ps_ty2 ty1 ty1 }
+    }
   where
     role = eqRelRole eq_rel
 
@@ -2268,15 +2280,15 @@ Main purpose: create new evidence for new_pred;
         Given           Already in inert               Nothing
                         Not                            Just new_evidence
 
-Note [Rewriting with Refl]
-~~~~~~~~~~~~~~~~~~~~~~~~~~
+Note [Rewriting with ZonkCo]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 If the coercion is just reflexivity then you may re-use the same
-variable.  But be careful!  Although the coercion is Refl, new_pred
+variable.  But be careful!  Although the coercion is ZonkCo, new_pred
 may reflect the result of unification alpha := ty, so new_pred might
 not _look_ the same as old_pred, and it's vital to proceed from now on
 using new_pred.
 
-qThe flattener preserves type synonyms, so they should appear in new_pred
+The flattener preserves type synonyms, so they should appear in new_pred
 as well as in old_pred; that is important for good error messages.
  -}
 
@@ -2348,8 +2360,8 @@ rewriteEqEvidence old_ev swapped nlhs nrhs lhs_co rhs_co
   = return (old_ev { ctev_pred = new_pred })
 
   | NotSwapped <- swapped
-  , isTcReflCo lhs_co      -- See Note [Rewriting with Refl]
-  , isTcReflCo rhs_co
+  , isZonkCo lhs_co      -- See Note [Rewriting with ZonkCo]
+  , isZonkCo rhs_co
   = return (old_ev { ctev_pred = new_pred })
 
   | CtGiven { ctev_evar = old_evar } <- old_ev
diff --git a/compiler/typecheck/TcSMonad.hs b/compiler/typecheck/TcSMonad.hs
index 8a2f8a2aa6..fd4cd39f68 100644
--- a/compiler/typecheck/TcSMonad.hs
+++ b/compiler/typecheck/TcSMonad.hs
@@ -1625,8 +1625,18 @@ add_item ics item@(CFunEqCan { cc_fun = tc, cc_tyargs = tys })
   = ics { inert_funeqs = insertFunEq (inert_funeqs ics) tc tys item }
 
 add_item ics item@(CTyEqCan { cc_tyvar = tv, cc_ev = ev })
-  = ics { inert_eqs   = addTyEq (inert_eqs ics) tv item
+  = ics { inert_eqs   = addTyEq (inert_eqs ics) tv item'
         , inert_count = bumpUnsolvedCount ev (inert_count ics) }
+  where
+    item' = case ev of
+              CtGiven { ctev_pred = pred, ctev_evar = co_var }
+                -> item { cc_ev = ev { ctev_evar = co_var `setCoVarPred` pred } }
+              _ -> item
+      -- The evidence Id in a Given equality in InertEqs must match
+      -- the predicate (which is not true of all CTyEqCans), else when
+      -- we fetch it out in TcFlatten.flatten_tyvar2, ctEvCoercion won't have the
+      -- correct type, and we'll mess up the WKTI
+      -- See Note [Ct/evidence invariant]
 
 add_item ics@(IC { inert_irreds = irreds, inert_count = count })
          item@(CIrredCan { cc_ev = ev, cc_insol = insoluble })
diff --git a/compiler/types/Coercion.hs b/compiler/types/Coercion.hs
index 8d07a681a1..b3b51f13c5 100644
--- a/compiler/types/Coercion.hs
+++ b/compiler/types/Coercion.hs
@@ -63,11 +63,11 @@ module Coercion (
 
         pickLR,
 
-        isGReflCo, isReflCo, isReflCo_maybe, isGReflCo_maybe, isReflexiveCo,
+        isGReflCo, isReflCo, isReflCo_maybe, isZonkCo, isGReflCo_maybe, isReflexiveCo,
         isReflCoVar_maybe, isGReflMCo, coToMCo,
 
         -- ** Coercion variables
-        mkCoVar, isCoVar, coVarName, setCoVarName, setCoVarUnique,
+        mkCoVar, isCoVar, coVarName, setCoVarName, setCoVarUnique, setCoVarPred,
         isCoVar_maybe,
 
         -- ** Free variables
@@ -169,7 +169,10 @@ setCoVarUnique :: CoVar -> Unique -> CoVar
 setCoVarUnique = setVarUnique
 
 setCoVarName :: CoVar -> Name -> CoVar
-setCoVarName   = setVarName
+setCoVarName = setVarName
+
+setCoVarPred :: CoVar -> PredType -> CoVar
+setCoVarPred = setVarType
 
 {-
 %************************************************************************
@@ -556,11 +559,28 @@ isGReflMCo _ = False
 -- | Tests if this coercion is obviously reflexive. Guaranteed to work
 -- very quickly. Sometimes a coercion can be reflexive, but not obviously
 -- so. c.f. 'isReflexiveCo'
+--
+--    isReflCo co  =>   co : t ~ t
 isReflCo :: Coercion -> Bool
 isReflCo (Refl{}) = True
 isReflCo (GRefl _ _ mco) | isGReflMCo mco = True
 isReflCo _ = False
 
+-- | Tests if this coercion is reflexive after zonkign.
+-- Guaranteed to work quickly. c.f. isReflCo
+--
+--    isZonkCo co  =>   co : t1 ~ t2, and zonk(t1)=zonk(t2)
+--
+-- ToDo: Is it better to have more cases here (eg TyConAppCo)
+--       or to make mkTyConAppCo smarter, to bubble up ZonkCos?
+isZonkCo :: Coercion -> Bool
+isZonkCo (ZonkCo {})          = True
+isZonkCo co | isReflCo co     = True
+isZonkCo (TyConAppCo _ _ cos) = all isZonkCo cos
+isZonkCo (AppCo co1 co2)      = isZonkCo co1 && isZonkCo co2
+isZonkCo _                    = False  -- Always safe
+
+
 -- | Returns the type coerced if this coercion is a generalized reflexive
 -- coercion. Guaranteed to work very quickly.
 isGReflCo_maybe :: Coercion -> Maybe (Type, Role)
diff --git a/testsuite/tests/polykinds/T14172.hs b/testsuite/tests/polykinds/T14172.hs
index 10fff5af69..a53ad5d98d 100644
--- a/testsuite/tests/polykinds/T14172.hs
+++ b/testsuite/tests/polykinds/T14172.hs
@@ -5,3 +5,6 @@ import T14172a
 
 traverseCompose :: (a -> f b) -> g a -> f (h _)
 traverseCompose = _Wrapping Compose . traverse
+
+-- newtype Compose f g a = Compose { getCompose :: f (g a) }
+-- Compose :: (k1 -> *) -> (k2 -> k1) -> k2 -> *
\ No newline at end of file
diff --git a/testsuite/tests/roles/should_compile/T8958.stderr b/testsuite/tests/roles/should_compile/T8958.stderr
index b975c66d9c..3ba5a968d0 100644
--- a/testsuite/tests/roles/should_compile/T8958.stderr
+++ b/testsuite/tests/roles/should_compile/T8958.stderr
@@ -16,58 +16,34 @@ CLASS INSTANCES
     -- Defined at T8958.hs:10:10
   instance [incoherent] Nominal a -- Defined at T8958.hs:7:10
 Dependent modules: []
-Dependent packages: [base-4.13.0.0, ghc-prim-0.6.1,
+Dependent packages: [base-4.14.0.0, ghc-prim-0.6.1,
                      integer-gmp-1.0.2.0]
 
 ==================== Typechecker ====================
 T8958.$tcMap
   = GHC.Types.TyCon
-      16542473435673943392##
-      5374201132143305512##
-      T8958.$trModule
-      (GHC.Types.TrNameS "Map"#)
-      0
-      GHC.Types.krep$*->*->*
+      16542473435673943392## 5374201132143305512## T8958.$trModule
+      (GHC.Types.TrNameS "Map"#) 0 GHC.Types.krep$*->*->*
 T8958.$tc'MkMap
   = GHC.Types.TyCon
-      2942839876828444488##
-      3989137838066763457##
-      T8958.$trModule
-      (GHC.Types.TrNameS "'MkMap"#)
-      2
-      $krep
+      2942839876828444488## 3989137838066763457## T8958.$trModule
+      (GHC.Types.TrNameS "'MkMap"#) 2 $krep
 T8958.$tcRepresentational
   = GHC.Types.TyCon
-      12809567151893673426##
-      12159693688248149156##
-      T8958.$trModule
-      (GHC.Types.TrNameS "Representational"#)
-      0
-      $krep
+      12809567151893673426## 12159693688248149156## T8958.$trModule
+      (GHC.Types.TrNameS "Representational"#) 0 $krep
 T8958.$tc'C:Representational
   = GHC.Types.TyCon
-      2358772282532242424##
-      5444038897914446879##
-      T8958.$trModule
-      (GHC.Types.TrNameS "'C:Representational"#)
-      1
-      $krep
+      2358772282532242424## 5444038897914446879## T8958.$trModule
+      (GHC.Types.TrNameS "'C:Representational"#) 1 $krep
 T8958.$tcNominal
   = GHC.Types.TyCon
-      12224997609886144634##
-      9866011944332051160##
-      T8958.$trModule
-      (GHC.Types.TrNameS "Nominal"#)
-      0
-      $krep
+      12224997609886144634## 9866011944332051160## T8958.$trModule
+      (GHC.Types.TrNameS "Nominal"#) 0 $krep
 T8958.$tc'C:Nominal
   = GHC.Types.TyCon
-      10562260635335201742##
-      1215478186250709459##
-      T8958.$trModule
-      (GHC.Types.TrNameS "'C:Nominal"#)
-      1
-      $krep
+      10562260635335201742## 1215478186250709459## T8958.$trModule
+      (GHC.Types.TrNameS "'C:Nominal"#) 1 $krep
 $krep [InlPrag=NOUSERINLINE[~]] = GHC.Types.KindRepVar 0
 $krep [InlPrag=NOUSERINLINE[~]] = GHC.Types.KindRepVar 1
 $krep [InlPrag=NOUSERINLINE[~]] = GHC.Types.KindRepFun $krep $krep
@@ -102,20 +78,24 @@ T8958.$trModule
   = GHC.Types.Module
       (GHC.Types.TrNameS "main"#) (GHC.Types.TrNameS "T8958"#)
 AbsBinds [a] []
-  {Exports: [T8958.$fRepresentationala <= $dRepresentational
-               wrap: <>]
-   Exported types: T8958.$fRepresentationala [InlPrag=NOUSERINLINE CONLIKE]
-                     :: forall a. Representational a
-                   [LclIdX[DFunId],
-                    Unf=DFun: \ (@ a) -> T8958.C:Representational TYPE: a]
+  {Exports: ABE { poly: T8958.$fRepresentationala [InlPrag=NOUSERINLINE CONLIKE]
+                          :: forall a. Representational a
+                        [LclIdX[DFunId],
+                         Unf=DFun: \ (@ a) -> T8958.C:Representational TYPE: a]
+                  mono: $dRepresentational :: Representational a
+                        [LclId]
+                  prags:
+                  wrap: <> }
    Binds: $dRepresentational = T8958.C:Representational @ a
    Evidence: [EvBinds{}]}
 AbsBinds [a] []
-  {Exports: [T8958.$fNominala <= $dNominal
-               wrap: <>]
-   Exported types: T8958.$fNominala [InlPrag=NOUSERINLINE CONLIKE]
-                     :: forall a. Nominal a
-                   [LclIdX[DFunId], Unf=DFun: \ (@ a) -> T8958.C:Nominal TYPE: a]
+  {Exports: ABE { poly: T8958.$fNominala [InlPrag=NOUSERINLINE CONLIKE]
+                          :: forall a. Nominal a
+                        [LclIdX[DFunId], Unf=DFun: \ (@ a) -> T8958.C:Nominal TYPE: a]
+                  mono: $dNominal :: Nominal a
+                        [LclId]
+                  prags:
+                  wrap: <> }
    Binds: $dNominal = T8958.C:Nominal @ a
    Evidence: [EvBinds{}]}