Fix kind-checking for data/newtypeswip/T16828

In one spot in kcConDecl we were passing in the return
kind signature rether than the return kind. e.g. #16828

   newtype instance Foo :: Type -> Type where
     MkFoo :: a -> Foo a

We were giving kcConDecl the kind (Type -> Type), whereas it
was expecting the ultimate return kind, namely Type.

This "looking past arrows" was being done, independently,
in several places, but we'd missed one.  This patch moves it all
to one place -- the new function kcConDecls (note the plural).

I also took the opportunity to rename
  tcDataFamHeader  to   tcDataFamInstHeader

(The previous name was consistently a source of confusion.)

4 files changed, 83 insertions, 31 deletions

diff --git a/compiler/typecheck/TcInstDcls.hs b/compiler/typecheck/TcInstDcls.hs
index c2f7a1100a..f4b8559abc 100644
--- a/compiler/typecheck/TcInstDcls.hs
+++ b/compiler/typecheck/TcInstDcls.hs
@@ -660,9 +660,9 @@ tcDataFamInstDecl mb_clsinfo
           -- Do /not/ check that the number of patterns = tyConArity fam_tc
           -- See [Arity of data families] in FamInstEnv
        ; (qtvs, pats, res_kind, stupid_theta)
-             <- tcDataFamHeader mb_clsinfo fam_tc imp_vars mb_bndrs
-                                fixity hs_ctxt hs_pats m_ksig hs_cons
-                                new_or_data
+             <- tcDataFamInstHeader mb_clsinfo fam_tc imp_vars mb_bndrs
+                                    fixity hs_ctxt hs_pats m_ksig hs_cons
+                                    new_or_data
 
        -- Eta-reduce the axiom if possible
        -- Quite tricky: see Note [Eta-reduction for data families]
@@ -782,16 +782,18 @@ tcDataFamInstDecl mb_clsinfo
 tcDataFamInstDecl _ _ = panic "tcDataFamInstDecl"
 
 -----------------------
-tcDataFamHeader :: AssocInstInfo -> TyCon -> [Name] -> Maybe [LHsTyVarBndr GhcRn]
-                -> LexicalFixity -> LHsContext GhcRn
-                -> HsTyPats GhcRn -> Maybe (LHsKind GhcRn) -> [LConDecl GhcRn]
-                -> NewOrData
-                -> TcM ([TyVar], [Type], Kind, ThetaType)
--- The "header" is the part other than the data constructors themselves
--- e.g.  data instance D [a] :: * -> * where ...
+tcDataFamInstHeader
+    :: AssocInstInfo -> TyCon -> [Name] -> Maybe [LHsTyVarBndr GhcRn]
+    -> LexicalFixity -> LHsContext GhcRn
+    -> HsTyPats GhcRn -> Maybe (LHsKind GhcRn) -> [LConDecl GhcRn]
+    -> NewOrData
+    -> TcM ([TyVar], [Type], Kind, ThetaType)
+-- The "header" of a data family instance is the part other than
+-- the data constructors themselves
+--    e.g.  data instance D [a] :: * -> * where ...
 -- Here the "header" is the bit before the "where"
-tcDataFamHeader mb_clsinfo fam_tc imp_vars mb_bndrs fixity hs_ctxt
-                hs_pats m_ksig hs_cons new_or_data
+tcDataFamInstHeader mb_clsinfo fam_tc imp_vars mb_bndrs fixity
+                    hs_ctxt hs_pats m_ksig hs_cons new_or_data
   = do { (imp_tvs, (exp_tvs, (stupid_theta, lhs_ty)))
             <- pushTcLevelM_                                $
                solveEqualities                              $
@@ -806,8 +808,10 @@ tcDataFamHeader mb_clsinfo fam_tc imp_vars mb_bndrs fixity hs_ctxt
 
                   -- Add constraints from the result signature
                   ; res_kind <- tc_kind_sig m_ksig
+
                   -- Add constraints from the data constructors
-                  ; mapM_ (wrapLocM_ (kcConDecl new_or_data res_kind)) hs_cons
+                  ; kcConDecls new_or_data res_kind hs_cons
+
                   ; lhs_ty <- checkExpectedKind_pp pp_lhs lhs_ty lhs_kind res_kind
                   ; return (stupid_theta, lhs_ty) }
 
@@ -832,7 +836,7 @@ tcDataFamHeader mb_clsinfo fam_tc imp_vars mb_bndrs fixity hs_ctxt
        -- the body of unravelFamInstPats.
        ; pats <- case splitTyConApp_maybe lhs_ty of
            Just (_, pats) -> pure pats
-           Nothing -> pprPanic "tcDataFamHeader" (ppr lhs_ty)
+           Nothing -> pprPanic "tcDataFamInstHeader" (ppr lhs_ty)
        ; return (qtvs, pats, typeKind lhs_ty, stupid_theta) }
           -- See Note [Unifying data family kinds] about why we need typeKind here
   where
@@ -889,7 +893,7 @@ is always redundant (except, perhaps, in that it helps guide unification). We
 have a definitive kind for the data family from the data family declaration,
 and so we learn nothing really new from the kind signature on an instance.
 We still must perform this unification (done in the call to checkExpectedKind
-toward the beginning of tcDataFamHeader), but the result is unhelpful. If there
+toward the beginning of tcDataFamInstHeader), but the result is unhelpful. If there
 is a cast, it will wrap the lhs_ty, and so we just drop it before splitting the
 lhs_ty to reveal the underlying patterns. Because of the potential of dropping
 a cast like this, we just use typeKind in the result instead of propagating res_kind
diff --git a/compiler/typecheck/TcTyClsDecls.hs b/compiler/typecheck/TcTyClsDecls.hs
index 06a730519b..0e49afd9b0 100644
--- a/compiler/typecheck/TcTyClsDecls.hs
+++ b/compiler/typecheck/TcTyClsDecls.hs
@@ -15,7 +15,7 @@ module TcTyClsDecls (
 
         -- Functions used by TcInstDcls to check
         -- data/type family instance declarations
-        kcConDecl, tcConDecls, dataDeclChecks, checkValidTyCon,
+        kcConDecls, tcConDecls, dataDeclChecks, checkValidTyCon,
         tcFamTyPats, tcTyFamInstEqn,
         tcAddTyFamInstCtxt, tcMkDataFamInstCtxt, tcAddDataFamInstCtxt,
         unravelFamInstPats, addConsistencyConstraints,
@@ -1138,8 +1138,7 @@ kcTyClDecl (DataDecl { tcdLName    = (dL->L _ name)
                , dd_ND = new_or_data } <- defn
   = do { tyCon <- kcLookupTcTyCon name
          -- See Note [Implementation of UnliftedNewtypes] STEP 2
-       ; (_, final_res_kind) <- etaExpandAlgTyCon (tyConBinders tyCon) (tyConResKind tyCon)
-       ; mapM_ (wrapLocM_ (kcConDecl new_or_data final_res_kind)) cons
+       ; kcConDecls new_or_data (tyConResKind tyCon) cons
        }
     -- hs_tvs and dd_kindSig already dealt with in getInitialKind
     -- This must be a GADT-style decl,
@@ -1154,7 +1153,7 @@ kcTyClDecl (DataDecl { tcdLName    = (dL->L _ name)
   = bindTyClTyVars name $ \ _ _ ->
     do { _ <- tcHsContext ctxt
        ; tyCon <- kcLookupTcTyCon name
-       ; mapM_ (wrapLocM_ (kcConDecl new_or_data (tyConResKind tyCon))) cons
+       ; kcConDecls new_or_data (tyConResKind tyCon) cons
        }
 
 kcTyClDecl (SynDecl { tcdLName = dL->L _ name, tcdRhs = rhs })
@@ -1223,16 +1222,27 @@ kcConArgTys new_or_data res_kind arg_tys = do
       unifyNewtypeKind dflags new_or_data arg_tys arg_tc_tys res_kind
   }
 
+kcConDecls :: NewOrData
+           -> Kind             -- The result kind signature
+           -> [LConDecl GhcRn] -- The data constructors
+           -> TcM ()
+kcConDecls new_or_data res_kind cons
+  = mapM_ (wrapLocM_ (kcConDecl new_or_data final_res_kind)) cons
+  where
+    (_, final_res_kind) = splitPiTys res_kind
+        -- See Note [kcConDecls result kind]
+
 -- Kind check a data constructor. In additional to the data constructor,
 -- we also need to know about whether or not its corresponding type was
 -- declared with data or newtype, and we need to know the result kind of
 -- this type. See Note [Implementation of UnliftedNewtypes] for why
 -- we need the first two arguments.
-kcConDecl ::
-     NewOrData -- Was the corresponding type declared with data or newtype?
-  -> Kind -- The result kind of the corresponding type constructor
-  -> ConDecl GhcRn -- The data constructor
-  -> TcM ()
+kcConDecl :: NewOrData
+          -> Kind  -- Result kind of the type constructor
+                   -- Usually Type but can be TYPE UnliftedRep
+                   -- or even TYPE r, in the case of unlifted newtype
+          -> ConDecl GhcRn
+          -> TcM ()
 kcConDecl new_or_data res_kind (ConDeclH98
   { con_name = name, con_ex_tvs = ex_tvs
   , con_mb_cxt = ex_ctxt, con_args = args })
@@ -1241,8 +1251,8 @@ kcConDecl new_or_data res_kind (ConDeclH98
     bindExplicitTKBndrs_Tv ex_tvs $
     do { _ <- tcHsMbContext ex_ctxt
        ; kcConArgTys new_or_data res_kind (hsConDeclArgTys args)
-         -- We don't need to check the telescope here, because that's
-         -- done in tcConDecl
+         -- We don't need to check the telescope here,
+         -- because that's done in tcConDecl
        }
 
 kcConDecl new_or_data res_kind (ConDeclGADT
@@ -1250,8 +1260,8 @@ kcConDecl new_or_data res_kind (ConDeclGADT
     , con_args = args, con_res_ty = res_ty })
   | HsQTvs { hsq_ext = implicit_tkv_nms
            , hsq_explicit = explicit_tkv_nms } <- qtvs
-  = -- Even though the data constructor's type is closed, we
-    -- must still kind-check the type, because that may influence
+  = -- Even though the GADT-style data constructor's type is closed,
+    -- we must still kind-check the type, because that may influence
     -- the inferred kind of the /type/ constructor.  Example:
     --    data T f a where
     --      MkT :: f a -> T f a
@@ -1269,7 +1279,31 @@ kcConDecl new_or_data res_kind (ConDeclGADT
 kcConDecl _ _ (ConDeclGADT _ _ _ (XLHsQTyVars _) _ _ _ _) = panic "kcConDecl"
 kcConDecl _ _ (XConDecl _) = panic "kcConDecl"
 
-{-
+{- Note [kcConDecls result kind]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+We might have e.g.
+    data T a :: Type -> Type where ...
+or
+    newtype instance N a :: Type -> Type  where ..
+in which case, the 'res_kind' passed to kcConDecls will be
+   Type->Type
+
+We must look past those arrows, or even foralls, to the Type in the
+corner, to pass to kcConDecl c.f. #16828. Hence the splitPiTys here.
+
+I am a bit concerned about tycons with a declaration like
+   data T a :: Type -> forall k. k -> Type  where ...
+
+It does not have a CUSK, so kcLHsQTyVars_NonCusk will make a TcTyCon
+with tyConResKind of Type -> forall k. k -> Type.  Even that is fine:
+the splitPiTys will look past the forall.  But I'm bothered about
+what if the type "in the corner" metions k?  This is incredibly
+obscure but something like this could be bad:
+   data T a :: Type -> foral k. k -> TYPE (F k) where ...
+
+I bet we are not quite right here, but my brain suffered a buffer
+overflow and I thought it best to nail the common cases right now.
+
 Note [Recursion and promoting data constructors]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We don't want to allow promotion in a strongly connected component
@@ -1493,7 +1527,7 @@ the H98 syntax doesn't permit a kind signature on the newtype itself.
 2. In a newtype instance (with -XUnliftedNewtypes), if the user does
    not write a kind signature, we want to allow the possibility that
    the kind is not Type, so we use newOpenTypeKind instead of liftedTypeKind.
-   This is done in tcDataFamHeader in TcInstDcls. Example:
+   This is done in tcDataFamInstHeader in TcInstDcls. Example:
 
        data family Bar (a :: RuntimeRep) :: TYPE a
        newtype instance Bar 'IntRep = BarIntC Int#
@@ -2246,7 +2280,7 @@ tcTyFamInstEqnGuts fam_tc mb_clsinfo imp_vars exp_bndrs hs_pats hs_rhs_ty
 
        -- See Note [Generalising in tcTyFamInstEqnGuts]
        -- This code (and the stuff immediately above) is very similar
-       -- to that in tcDataFamHeader.  Maybe we should abstract the
+       -- to that in tcDataFamInstHeader.  Maybe we should abstract the
        -- common code; but for the moment I concluded that it's
        -- clearer to duplicate it.  Still, if you fix a bug here,
        -- check there too!
diff --git a/testsuite/tests/indexed-types/should_compile/T16828.hs b/testsuite/tests/indexed-types/should_compile/T16828.hs
new file mode 100644
index 0000000000..23bfc25dcb
--- /dev/null
+++ b/testsuite/tests/indexed-types/should_compile/T16828.hs
@@ -0,0 +1,13 @@
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UnliftedNewtypes #-}
+
+module Bug where
+
+import Data.Kind
+
+data family Foo :: Type -> Type
+
+-- This declaration is fine
+newtype instance Foo :: Type -> Type where
+  MkFoo :: a -> Foo a
diff --git a/testsuite/tests/indexed-types/should_compile/all.T b/testsuite/tests/indexed-types/should_compile/all.T
index 9142b8edcd..4a15bdf4f4 100644
--- a/testsuite/tests/indexed-types/should_compile/all.T
+++ b/testsuite/tests/indexed-types/should_compile/all.T
@@ -290,3 +290,4 @@ test('T16110_Compile', normal, compile, [''])
 test('T16356_Compile1', normal, compile, [''])
 test('T16356_Compile2', normal, compile, [''])
 test('T16632', normal, compile, ['-Wunused-type-patterns -fdiagnostics-show-caret'])
+test('T16828', normal, compile, [''])