Refactor UnliftedNewtypes-relation kind signature validity checks

This fixes three infelicities related to the programs that are
(and aren't) accepted with `UnliftedNewtypes`:

* Enabling `UnliftedNewtypes` would permit newtypes to have return
  kind `Id Type`, which had disastrous results (i.e., GHC panics).
* Data family declarations ending in kind `TYPE r` (for some `r`)
  weren't being accepted if `UnliftedNewtypes` wasn't enabled,
  despite the GHC proposal specifying otherwise.
* GHC wasn't warning about programs that _would_ typecheck if
  `UnliftedNewtypes` were enabled in certain common cases.

As part of fixing these issues, I factored out the logic for checking
all of the various properties about data type/data family return
kinds into a single `checkDataKindSig` function. I also cleaned up
some of the formatting in the existing error message that gets
thrown.

Fixes #16821, fixes #16827, and fixes #16829.

17 files changed, 220 insertions, 67 deletions

diff --git a/compiler/typecheck/TcHsType.hs b/compiler/typecheck/TcHsType.hs
index adc25f5ff8..fac6f2e31c 100644
--- a/compiler/typecheck/TcHsType.hs
+++ b/compiler/typecheck/TcHsType.hs
@@ -49,7 +49,7 @@ module TcHsType (
         kindGeneralize, checkExpectedKind_pp,
 
         -- Sort-checking kinds
-        tcLHsKindSig, badKindSig,
+        tcLHsKindSig, checkDataKindSig, DataSort(..),
 
         -- Zonking and promoting
         zonkPromoteType,
@@ -105,7 +105,7 @@ import UniqSupply
 import Outputable
 import FastString
 import PrelNames hiding ( wildCardName )
-import DynFlags ( WarningFlag (Opt_WarnPartialTypeSignatures) )
+import DynFlags
 import qualified GHC.LanguageExtensions as LangExt
 
 import Maybes
@@ -2405,12 +2405,101 @@ etaExpandAlgTyCon tc_bndrs kind
               (subst', tv') = substTyVarBndr subst tv
               tcb = Bndr tv' (NamedTCB vis)
 
-badKindSig :: Bool -> Kind -> SDoc
-badKindSig check_for_type kind
- = hang (sep [ text "Kind signature on data type declaration has non-*"
-             , (if check_for_type then empty else text "and non-variable") <+>
-               text "return kind" ])
-        2 (ppr kind)
+-- | A description of whether something is a
+--
+-- * @data@ or @newtype@ ('DataDeclSort')
+--
+-- * @data instance@ or @newtype instance@ ('DataInstanceSort')
+--
+-- * @data family@ ('DataFamilySort')
+--
+-- At present, this data type is only consumed by 'checkDataKindSig'.
+data DataSort
+  = DataDeclSort     NewOrData
+  | DataInstanceSort NewOrData
+  | DataFamilySort
+
+-- | Checks that the return kind in a data declaration's kind signature is
+-- permissible. There are three cases:
+--
+-- If dealing with a @data@, @newtype@, @data instance@, or @newtype instance@
+-- declaration, check that the return kind is @Type@.
+--
+-- If the declaration is a @newtype@ or @newtype instance@ and the
+-- @UnliftedNewtypes@ extension is enabled, this check is slightly relaxed so
+-- that a return kind of the form @TYPE r@ (for some @r@) is permitted.
+-- See @Note [Implementation of UnliftedNewtypes]@ in "TcTyClsDecls".
+--
+-- If dealing with a @data family@ declaration, check that the return kind is
+-- either of the form:
+--
+-- 1. @TYPE r@ (for some @r@), or
+--
+-- 2. @k@ (where @k@ is a bare kind variable; see #12369)
+checkDataKindSig :: DataSort -> Kind -> TcM ()
+checkDataKindSig data_sort kind = do
+  dflags <- getDynFlags
+  checkTc (is_TYPE_or_Type dflags || is_kind_var) (err_msg dflags)
+  where
+    pp_dec :: SDoc
+    pp_dec = text $
+      case data_sort of
+        DataDeclSort     DataType -> "data type"
+        DataDeclSort     NewType  -> "newtype"
+        DataInstanceSort DataType -> "data instance"
+        DataInstanceSort NewType  -> "newtype instance"
+        DataFamilySort            -> "data family"
+
+    is_newtype :: Bool
+    is_newtype =
+      case data_sort of
+        DataDeclSort     new_or_data -> new_or_data == NewType
+        DataInstanceSort new_or_data -> new_or_data == NewType
+        DataFamilySort               -> False
+
+    is_data_family :: Bool
+    is_data_family =
+      case data_sort of
+        DataDeclSort{}     -> False
+        DataInstanceSort{} -> False
+        DataFamilySort     -> True
+
+    tYPE_ok :: DynFlags -> Bool
+    tYPE_ok dflags =
+         (is_newtype && xopt LangExt.UnliftedNewtypes dflags)
+           -- With UnliftedNewtypes, we allow kinds other than Type, but they
+           -- must still be of the form `TYPE r` since we don't want to accept
+           -- Constraint or Nat.
+           -- See Note [Implementation of UnliftedNewtypes] in TcTyClsDecls.
+      || is_data_family
+           -- If this is a `data family` declaration, we don't need to check if
+           -- UnliftedNewtypes is enabled, since data family declarations can
+           -- have return kind `TYPE r` unconditionally (#16827).
+
+    is_TYPE :: Bool
+    is_TYPE = tcIsRuntimeTypeKind kind
+
+    is_TYPE_or_Type :: DynFlags -> Bool
+    is_TYPE_or_Type dflags | tYPE_ok dflags = is_TYPE
+                           | otherwise      = tcIsLiftedTypeKind kind
+
+    -- In the particular case of a data family, permit a return kind of the
+    -- form `:: k` (where `k` is a bare kind variable).
+    is_kind_var :: Bool
+    is_kind_var | is_data_family = isJust (tcGetCastedTyVar_maybe kind)
+                | otherwise      = False
+
+    err_msg :: DynFlags -> SDoc
+    err_msg dflags =
+      sep [ (sep [ text "Kind signature on" <+> pp_dec <+>
+                   text "declaration has non-" <>
+                   (if tYPE_ok dflags then text "TYPE" else ppr liftedTypeKind)
+                 , (if is_data_family then text "and non-variable" else empty) <+>
+                   text "return kind" <+> quotes (ppr kind) ])
+          , if not (tYPE_ok dflags) && is_TYPE && is_newtype &&
+               not (xopt LangExt.UnliftedNewtypes dflags)
+            then text "Perhaps you intended to use UnliftedNewtypes"
+            else empty ]
 
 tcbVisibilities :: TyCon -> [Type] -> [TyConBndrVis]
 -- Result is in 1-1 correpondence with orig_args
diff --git a/compiler/typecheck/TcInstDcls.hs b/compiler/typecheck/TcInstDcls.hs
index 25c598df6b..c2f7a1100a 100644
--- a/compiler/typecheck/TcInstDcls.hs
+++ b/compiler/typecheck/TcInstDcls.hs
@@ -687,17 +687,7 @@ tcDataFamInstDecl mb_clsinfo
        --     we did it before the "extra" tvs from etaExpandAlgTyCon
        --     would always be eta-reduced
        ; (extra_tcbs, final_res_kind) <- etaExpandAlgTyCon full_tcbs res_kind
-       ; unlifted_newtypes <- xoptM LangExt.UnliftedNewtypes
-       ; let allowUnlifted = unlifted_newtypes && new_or_data == NewType
-         -- With UnliftedNewtypes, we allow kinds other than Type, but they
-         -- must still be of the form `TYPE r` since we don't want to accept
-         -- Constraint or Nat. See Note [Implementation of UnliftedNewtypes].
-       ; checkTc
-           (if allowUnlifted
-              then tcIsRuntimeTypeKind final_res_kind
-              else tcIsLiftedTypeKind final_res_kind
-           )
-           (badKindSig True res_kind)
+       ; checkDataKindSig (DataInstanceSort new_or_data) final_res_kind
        ; let extra_pats  = map (mkTyVarTy . binderVar) extra_tcbs
              all_pats    = pats `chkAppend` extra_pats
              orig_res_ty = mkTyConApp fam_tc all_pats
diff --git a/compiler/typecheck/TcTyClsDecls.hs b/compiler/typecheck/TcTyClsDecls.hs
index 9d2aea8d15..06a730519b 100644
--- a/compiler/typecheck/TcTyClsDecls.hs
+++ b/compiler/typecheck/TcTyClsDecls.hs
@@ -1877,15 +1877,7 @@ tcFamDecl1 parent (FamilyDecl { fdInfo = fam_info
   -- When UnliftedNewtypes is enabled, we loosen this restriction
   -- on the return kind. See Note [Implementation of UnliftedNewtypes], wrinkle (1).
   ; let (_, final_res_kind) = splitPiTys res_kind
-  ; unlifted_newtypes <- xoptM LangExt.UnliftedNewtypes
-  ; checkTc
-      (  (if unlifted_newtypes
-            then tcIsRuntimeTypeKind final_res_kind
-            else tcIsLiftedTypeKind final_res_kind
-         )
-      || isJust (tcGetCastedTyVar_maybe final_res_kind)
-      )
-      (badKindSig False res_kind)
+  ; checkDataKindSig DataFamilySort final_res_kind
   ; tc_rep_name <- newTyConRepName tc_name
   ; let tycon = mkFamilyTyCon tc_name binders
                               res_kind
@@ -2033,15 +2025,8 @@ tcDataDefn roles_info
        ; tcg_env <- getGblEnv
        ; (extra_bndrs, final_res_kind) <- etaExpandAlgTyCon tycon_binders res_kind
        ; let hsc_src = tcg_src tcg_env
-       ; unlifted_newtypes <- xoptM LangExt.UnliftedNewtypes
-       ; let allowUnlifted = unlifted_newtypes && new_or_data == NewType
-       ; unless (mk_permissive_kind hsc_src cons || allowUnlifted) $
-           checkTc
-             (if allowUnlifted
-                then tcIsRuntimeTypeKind final_res_kind
-                else tcIsLiftedTypeKind final_res_kind
-             )
-             (badKindSig True res_kind)
+       ; unless (mk_permissive_kind hsc_src cons) $
+           checkDataKindSig (DataDeclSort new_or_data) final_res_kind
 
        ; stupid_tc_theta <- pushTcLevelM_ $ solveEqualities $ tcHsContext ctxt
        ; stupid_theta    <- zonkTcTypesToTypes stupid_tc_theta
@@ -2077,7 +2062,12 @@ tcDataDefn roles_info
   where
     -- Abstract data types in hsig files can have arbitrary kinds,
     -- because they may be implemented by type synonyms
-    -- (which themselves can have arbitrary kinds, not just *)
+    -- (which themselves can have arbitrary kinds, not just *). See #13955.
+    --
+    -- Note that this is only a property that data type declarations possess,
+    -- so one could not have, say, a data family instance in an hsig file that
+    -- has kind `Bool`. Therfore, this check need only occur in the code that
+    -- typechecks data type declarations.
     mk_permissive_kind HsigFile [] = True
     mk_permissive_kind _ _ = False
 
@@ -2600,7 +2590,7 @@ tcConDecl rep_tycon tag_map tmpl_bndrs _res_kind res_tmpl new_or_data
                  ; btys <- tcConArgs hs_args
                  ; let (arg_tys, stricts) = unzip btys
                  ; res_ty <- tcHsOpenType hs_res_ty
-                   -- See Note [Implementation of unlifted newtypes]
+                   -- See Note [Implementation of UnliftedNewtypes]
                  ; dflags <- getDynFlags
                  ; final_arg_tys <-
                      unifyNewtypeKind dflags new_or_data
diff --git a/docs/users_guide/glasgow_exts.rst b/docs/users_guide/glasgow_exts.rst
index 32818e8c70..570548095b 100644
--- a/docs/users_guide/glasgow_exts.rst
+++ b/docs/users_guide/glasgow_exts.rst
@@ -179,7 +179,7 @@ There are some restrictions on the use of primitive types:
          newtype A = MkA Int#
 
    However, this restriction can be relaxed by enabling
-   :extension:`-XUnliftedNewtypes`.  The `section on unlifted newtypes
+   :extension:`UnliftedNewtypes`.  The `section on unlifted newtypes
    <#unlifted-newtypes>`__ details the behavior of such types.
 
 -  You cannot bind a variable with an unboxed type in a *top-level*
@@ -375,7 +375,9 @@ Unlifted Newtypes
 
     Enable the use of newtypes over types with non-lifted runtime representations.
 
-``-XUnliftedNewtypes`` relaxes the restrictions around what types can appear inside
+GHC implements an :extension:`UnliftedNewtypes` extension as specified in
+`this GHC proposal <https://github.com/ghc-proposals/ghc-proposals/blob/master/proposals/0013-unlifted-newtypes.rst>`_.
+:extension:`UnliftedNewtypes` relaxes the restrictions around what types can appear inside
 of a `newtype`. For example, the type ::
 
     newtype A = MkA Int#
@@ -410,28 +412,28 @@ And with `UnboxedSums <#unboxed-sums>`__ enabled ::
     newtype Maybe# :: forall (r :: RuntimeRep). TYPE r -> TYPE (SumRep '[r, TupleRep '[]]) where
       MkMaybe# :: forall (r :: RuntimeRep) (a :: TYPE r). (# a | (# #) #) -> Maybe# a
 
-This extension also relaxes some of the restrictions around data families.
-It must be enabled in modules where either of the following occur:
-
-- A data family is declared with a kind other than ``Type``. Both ``Foo``
-  and ``Bar``, defined below, fall into this category:
-  ::
+This extension also relaxes some of the restrictions around data family
+instances. In particular, :extension:`UnliftedNewtypes` permits a
+``newtype instance`` to be given a return kind of ``TYPE r``, not just
+``Type``. For example, the following ``newtype instance`` declarations would be
+permitted: ::
 
      class Foo a where
        data FooKey a :: TYPE 'IntRep
      class Bar (r :: RuntimeRep) where
        data BarType r :: TYPE r
 
-- A ``newtype instance`` is written with a kind other than ``Type``. The
-  following instances of ``Foo`` and ``Bar`` as defined above fall into
-  this category.
-  ::
-
      instance Foo Bool where
        newtype FooKey Bool = FooKeyBoolC Int#
      instance Bar 'WordRep where
        newtype BarType 'WordRep = BarTypeWordRepC Word#
 
+It is worth noting that :extension:`UnliftedNewtypes` is *not* required to give
+the data families themselves return kinds involving ``TYPE``, such as the
+``FooKey`` and ``BarType`` examples above. The extension is
+only required for ``newtype instance`` declarations, such as ``FooKeyBoolC``
+and ``BarTypeWorkRepC`` above.
+
 This extension impacts the determination of whether or not a newtype has
 a Complete User-Specified Kind Signature (CUSK). The exact impact is specified
 `the section on CUSKs <#complete-kind-signatures>`__.
@@ -7650,9 +7652,26 @@ entirely optional, so that we can declare ``Array`` alternatively with ::
     data family Array :: Type -> Type
 
 Unlike with ordinary data definitions, the result kind of a data family
-does not need to be ``Type``: it can alternatively be a kind variable
-(with :extension:`PolyKinds`). Data instances' kinds must end in
-``Type``, however.
+does not need to be ``Type``. It can alternatively be:
+
+* Of the form ``TYPE r`` for some ``r`` (see :ref:`runtime-rep`).
+  For example: ::
+
+    data family DF1 :: TYPE IntRep
+    data family DF2 (r :: RuntimeRep)  :: TYPE r
+    data family DF3 :: Type -> TYPE WordRep
+
+* A bare kind variable (with :extension:`PolyKinds` enabled).
+  For example: ::
+
+    data family DF4 :: k
+    data family DF5 (a :: k) :: k
+    data family DF6 :: (k -> Type) -> k
+
+Data instances' kinds must end in ``Type``, however. This restriction is
+slightly relaxed when the :extension:`UnliftedNewtypes` extension is enabled,
+as it permits a ``newtype instance``'s kind to end in ``TYPE r`` for some
+``r``.
 
 .. _data-instance-declarations:
 
diff --git a/testsuite/tests/typecheck/should_compile/T16827.hs b/testsuite/tests/typecheck/should_compile/T16827.hs
new file mode 100644
index 0000000000..4209d6a592
--- /dev/null
+++ b/testsuite/tests/typecheck/should_compile/T16827.hs
@@ -0,0 +1,8 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE TypeFamilies #-}
+module T16827 where
+
+import Data.Kind
+import GHC.Exts
+
+data family Foo (a :: Type) :: TYPE 'IntRep
diff --git a/testsuite/tests/typecheck/should_compile/all.T b/testsuite/tests/typecheck/should_compile/all.T
index 8534a2c327..244aaa2ec4 100644
--- a/testsuite/tests/typecheck/should_compile/all.T
+++ b/testsuite/tests/typecheck/should_compile/all.T
@@ -673,6 +673,7 @@ test('T16225', normal, compile, [''])
 test('T13951', normal, compile, [''])
 test('T16411', normal, compile, [''])
 test('T16609', normal, compile, [''])
+test('T16827', normal, compile, [''])
 test('T505', normal, compile, [''])
 test('T12928', normal, compile, [''])
 test('UnliftedNewtypesGnd', normal, compile, [''])
diff --git a/testsuite/tests/typecheck/should_fail/T14048a.stderr b/testsuite/tests/typecheck/should_fail/T14048a.stderr
index 48a91c7525..d75b423df3 100644
--- a/testsuite/tests/typecheck/should_fail/T14048a.stderr
+++ b/testsuite/tests/typecheck/should_fail/T14048a.stderr
@@ -1,5 +1,5 @@
 
 T14048a.hs:6:1: error:
-    • Kind signature on data type declaration has non-* return kind
-        Constraint
+    • Kind signature on data type declaration has non-*
+      return kind ‘Constraint’
     • In the data declaration for ‘Foo’
diff --git a/testsuite/tests/typecheck/should_fail/T14048b.stderr b/testsuite/tests/typecheck/should_fail/T14048b.stderr
index fe78d9f7f5..4356d65ef7 100644
--- a/testsuite/tests/typecheck/should_fail/T14048b.stderr
+++ b/testsuite/tests/typecheck/should_fail/T14048b.stderr
@@ -1,6 +1,5 @@
 
 T14048b.hs:7:1: error:
-    • Kind signature on data type declaration has non-*
-      and non-variable return kind
-        Constraint
+    • Kind signature on data family declaration has non-TYPE
+      and non-variable return kind ‘Constraint’
     • In the data family declaration for ‘Foo’
diff --git a/testsuite/tests/typecheck/should_fail/T14048c.stderr b/testsuite/tests/typecheck/should_fail/T14048c.stderr
index 7e83d1924c..a0465fe4fc 100644
--- a/testsuite/tests/typecheck/should_fail/T14048c.stderr
+++ b/testsuite/tests/typecheck/should_fail/T14048c.stderr
@@ -1,5 +1,5 @@
 
 T14048c.hs:9:1: error:
-    • Kind signature on data type declaration has non-* return kind
-        Constraint
+    • Kind signature on data instance declaration has non-*
+      return kind ‘Constraint’
     • In the data instance declaration for ‘Foo’
diff --git a/testsuite/tests/typecheck/should_fail/T16821.hs b/testsuite/tests/typecheck/should_fail/T16821.hs
new file mode 100644
index 0000000000..2b23b90a7b
--- /dev/null
+++ b/testsuite/tests/typecheck/should_fail/T16821.hs
@@ -0,0 +1,16 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UnliftedNewtypes #-}
+module T16821 where
+
+import Data.Kind
+
+type family Id (x :: Type) :: Type where
+  Id x = x
+
+newtype T :: Id Type where
+  MkT :: Int -> T
+
+f :: T -> T
+f (MkT x) = MkT (x + 1)
diff --git a/testsuite/tests/typecheck/should_fail/T16821.stderr b/testsuite/tests/typecheck/should_fail/T16821.stderr
new file mode 100644
index 0000000000..f5d77720c7
--- /dev/null
+++ b/testsuite/tests/typecheck/should_fail/T16821.stderr
@@ -0,0 +1,5 @@
+
+T16821.hs:12:1: error:
+    • Kind signature on newtype declaration has non-TYPE
+      return kind ‘Id *’
+    • In the newtype declaration for ‘T’
diff --git a/testsuite/tests/typecheck/should_fail/T16829a.hs b/testsuite/tests/typecheck/should_fail/T16829a.hs
new file mode 100644
index 0000000000..20d8ace945
--- /dev/null
+++ b/testsuite/tests/typecheck/should_fail/T16829a.hs
@@ -0,0 +1,10 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE TypeFamilies #-}
+module T16829a where
+
+import GHC.Exts
+
+newtype T :: TYPE IntRep where
+  MkT :: Int# -> T
diff --git a/testsuite/tests/typecheck/should_fail/T16829a.stderr b/testsuite/tests/typecheck/should_fail/T16829a.stderr
new file mode 100644
index 0000000000..7ea8845cc9
--- /dev/null
+++ b/testsuite/tests/typecheck/should_fail/T16829a.stderr
@@ -0,0 +1,6 @@
+
+T16829a.hs:9:1: error:
+    • Kind signature on newtype declaration has non-*
+      return kind ‘TYPE 'IntRep’
+      Perhaps you intended to use UnliftedNewtypes
+    • In the newtype declaration for ‘T’
diff --git a/testsuite/tests/typecheck/should_fail/T16829b.hs b/testsuite/tests/typecheck/should_fail/T16829b.hs
new file mode 100644
index 0000000000..0375818d94
--- /dev/null
+++ b/testsuite/tests/typecheck/should_fail/T16829b.hs
@@ -0,0 +1,11 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE TypeFamilies #-}
+module T16829b where
+
+import GHC.Exts
+
+data family T :: TYPE IntRep
+newtype instance T :: TYPE IntRep where
+  MkT :: Int# -> T
diff --git a/testsuite/tests/typecheck/should_fail/T16829b.stderr b/testsuite/tests/typecheck/should_fail/T16829b.stderr
new file mode 100644
index 0000000000..590a884dc8
--- /dev/null
+++ b/testsuite/tests/typecheck/should_fail/T16829b.stderr
@@ -0,0 +1,6 @@
+
+T16829b.hs:10:1: error:
+    • Kind signature on newtype instance declaration has non-*
+      return kind ‘TYPE 'IntRep’
+      Perhaps you intended to use UnliftedNewtypes
+    • In the newtype instance declaration for ‘T’
diff --git a/testsuite/tests/typecheck/should_fail/UnliftedNewtypesConstraintFamily.stderr b/testsuite/tests/typecheck/should_fail/UnliftedNewtypesConstraintFamily.stderr
index 9c6816b3c1..93a48850dc 100644
--- a/testsuite/tests/typecheck/should_fail/UnliftedNewtypesConstraintFamily.stderr
+++ b/testsuite/tests/typecheck/should_fail/UnliftedNewtypesConstraintFamily.stderr
@@ -1,5 +1,5 @@
-UnliftedNewtypesConstraintFamily.hs:11:1:
-     Kind signature on data type declaration has non-*
-      and non-variable return kind
-        Constraint
-     In the data family declaration for ‘D’
+
+UnliftedNewtypesConstraintFamily.hs:11:1: error:
+    • Kind signature on data family declaration has non-TYPE
+      and non-variable return kind ‘Constraint’
+    • In the data family declaration for ‘D’
diff --git a/testsuite/tests/typecheck/should_fail/all.T b/testsuite/tests/typecheck/should_fail/all.T
index 0b0ae59391..fd6790bb46 100644
--- a/testsuite/tests/typecheck/should_fail/all.T
+++ b/testsuite/tests/typecheck/should_fail/all.T
@@ -524,6 +524,9 @@ test('T15883b', normal, compile_fail, [''])
 test('T15883c', normal, compile_fail, [''])
 test('T15883d', normal, compile_fail, [''])
 test('T15883e', normal, compile_fail, [''])
+test('T16821', normal, compile_fail, [''])
+test('T16829a', normal, compile_fail, [''])
+test('T16829b', normal, compile_fail, [''])
 test('UnliftedNewtypesFail', normal, compile_fail, [''])
 test('UnliftedNewtypesNotEnabled', normal, compile_fail, [''])
 test('UnliftedNewtypesCoerceFail', normal, compile_fail, [''])