Reify oversaturated data family instances correctly (#17296)wip/T17296

`TcSplice` was not properly handling oversaturated data family
instances, such as the example in #17296, as it dropped arguments due
to carelessly zipping data family instance arguments with
`tyConTyVars`. For data families, the number of `tyConTyVars` can
sometimes be less than the number of arguments it can accept in a
data family instance due to the fact that data family instances can
be oversaturated.

To account for this, `TcSplice.mkIsPolyTvs` has now been renamed to
`tyConArgsPolyKinded` and now factors in `tyConResKind` in addition
to `tyConTyVars`. I've also added
`Note [Reified instances and explicit kind signatures]` which
explains the various subtleties in play here.

Fixes #17296.

5 files changed, 175 insertions, 21 deletions

diff --git a/compiler/typecheck/TcSplice.hs b/compiler/typecheck/TcSplice.hs
index 05c2b0fd10..aac3a5da03 100644
--- a/compiler/typecheck/TcSplice.hs
+++ b/compiler/typecheck/TcSplice.hs
@@ -1478,13 +1478,11 @@ reifyAxBranch fam_tc (CoAxBranch { cab_tvs = tvs
   = do { tvs' <- reifyTyVarsToMaybe tvs
        ; let lhs_types_only = filterOutInvisibleTypes fam_tc lhs
        ; lhs' <- reifyTypes lhs_types_only
-       ; annot_th_lhs <- zipWith3M annotThType (mkIsPolyTvs fam_tvs)
+       ; annot_th_lhs <- zipWith3M annotThType (tyConArgsPolyKinded fam_tc)
                                    lhs_types_only lhs'
        ; let lhs_type = mkThAppTs (TH.ConT $ reifyName fam_tc) annot_th_lhs
        ; rhs'  <- reifyType rhs
        ; return (TH.TySynEqn tvs' lhs_type rhs') }
-  where
-    fam_tvs = tyConVisibleTyVars fam_tc
 
 reifyTyCon :: TyCon -> TcM TH.Info
 reifyTyCon tc
@@ -1713,7 +1711,8 @@ reifyClass cls
 -- | Annotate (with TH.SigT) a type if the first parameter is True
 -- and if the type contains a free variable.
 -- This is used to annotate type patterns for poly-kinded tyvars in
--- reifying class and type instances. See #8953 and th/T8953.
+-- reifying class and type instances.
+-- See @Note [Reified instances and explicit kind signatures]@.
 annotThType :: Bool   -- True <=> annotate
             -> TyCoRep.Type -> TH.Type -> TcM TH.Type
   -- tiny optimization: if the type is annotated, don't annotate again.
@@ -1725,24 +1724,116 @@ annotThType True ty th_ty
        ; return (TH.SigT th_ty th_ki) }
 annotThType _    _ th_ty = return th_ty
 
--- | For every type variable in the input,
--- report whether or not the tv is poly-kinded. This is used to eventually
--- feed into 'annotThType'.
-mkIsPolyTvs :: [TyVar] -> [Bool]
-mkIsPolyTvs = map is_poly_tv
+-- | For every argument type that a type constructor accepts,
+-- report whether or not the argument is poly-kinded. This is used to
+-- eventually feed into 'annotThType'.
+-- See @Note [Reified instances and explicit kind signatures]@.
+tyConArgsPolyKinded :: TyCon -> [Bool]
+tyConArgsPolyKinded tc =
+     map (is_poly_ty . tyVarKind)      tc_vis_tvs
+     -- See "Wrinkle: Oversaturated data family instances" in
+     -- @Note [Reified instances and explicit kind signatures]@
+  ++ map (is_poly_ty . tyCoBinderType) tc_res_kind_vis_bndrs -- (1) in Wrinkle
+  ++ repeat True                                             -- (2) in Wrinkle
   where
-    is_poly_tv tv = not $
+    is_poly_ty :: Type -> Bool
+    is_poly_ty ty = not $
                     isEmptyVarSet $
                     filterVarSet isTyVar $
-                    tyCoVarsOfType $
-                    tyVarKind tv
+                    tyCoVarsOfType ty
+
+    tc_vis_tvs :: [TyVar]
+    tc_vis_tvs = tyConVisibleTyVars tc
+
+    tc_res_kind_vis_bndrs :: [TyCoBinder]
+    tc_res_kind_vis_bndrs = filter isVisibleBinder $ fst $ splitPiTys $ tyConResKind tc
+
+{-
+Note [Reified instances and explicit kind signatures]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Reified class instances and type family instances often include extra kind
+information to disambiguate instances. Here is one such example that
+illustrates this (#8953):
+
+    type family Poly (a :: k) :: Type
+    type instance Poly (x :: Bool)    = Int
+    type instance Poly (x :: Maybe k) = Double
+
+If you're not careful, reifying these instances might yield this:
+
+    type instance Poly x = Int
+    type instance Poly x = Double
+
+To avoid this, we go through some care to annotate things with extra kind
+information. Some functions which accomplish this feat include:
+
+* annotThType: This annotates a type with a kind signature if the type contains
+  a free variable.
+* tyConArgsPolyKinded: This checks every argument that a type constructor can
+  accept and reports if the type of the argument is poly-kinded. This
+  information is ultimately fed into annotThType.
+
+-----
+-- Wrinkle: Oversaturated data family instances
+-----
+
+What constitutes an argument to a type constructor in the definition of
+tyConArgsPolyKinded? For most type constructors, it's simply the visible
+type variable binders (i.e., tyConVisibleTyVars). There is one corner case
+we must keep in mind, however: data family instances can appear oversaturated
+(#17296). For instance:
+
+    data family   Foo :: Type -> Type
+    data instance Foo x
+
+    data family Bar :: k
+    data family Bar x
+
+For these sorts of data family instances, tyConVisibleTyVars isn't enough,
+as they won't give you the kinds of the oversaturated arguments. We must
+also consult:
+
+1. The kinds of the arguments in the result kind (i.e., the tyConResKind).
+   This will tell us, e.g., the kind of `x` in `Foo x` above.
+2. If we go beyond the number of arguments in the result kind (like the
+   `x` in `Bar x`), then we conservatively assume that the argument's
+   kind is poly-kinded.
+
+-----
+-- Wrinkle: data family instances with return kinds
+-----
+
+Another squirrelly corner case is this:
+
+    data family Foo (a :: k)
+    data instance Foo :: Bool -> Type
+    data instance Foo :: Char -> Type
+
+If you're not careful, reifying these instances might yield this:
+
+    data instance Foo
+    data instance Foo
+
+We can fix this ambiguity by reifying the instances' explicit return kinds. We
+should only do this if necessary (see
+Note [When does a tycon application need an explicit kind signature?] in Type),
+but more importantly, we *only* do this if either of the following are true:
+
+1. The data family instance has no constructors.
+2. The data family instance is declared with GADT syntax.
+
+If neither of these are true, then reifying the return kind would yield
+something like this:
+
+    data instance (Bar a :: Type) = MkBar a
+
+Which is not valid syntax.
+-}
 
 ------------------------------
 reifyClassInstances :: Class -> [ClsInst] -> TcM [TH.Dec]
 reifyClassInstances cls insts
-  = mapM (reifyClassInstance (mkIsPolyTvs tvs)) insts
-  where
-    tvs = tyConVisibleTyVars (classTyCon cls)
+  = mapM (reifyClassInstance (tyConArgsPolyKinded (classTyCon cls))) insts
 
 reifyClassInstance :: [Bool]  -- True <=> the corresponding tv is poly-kinded
                               -- includes only *visible* tvs
@@ -1768,9 +1859,7 @@ reifyClassInstance is_poly_tvs i
 ------------------------------
 reifyFamilyInstances :: TyCon -> [FamInst] -> TcM [TH.Dec]
 reifyFamilyInstances fam_tc fam_insts
-  = mapM (reifyFamilyInstance (mkIsPolyTvs fam_tvs)) fam_insts
-  where
-    fam_tvs = tyConVisibleTyVars fam_tc
+  = mapM (reifyFamilyInstance (tyConArgsPolyKinded fam_tc)) fam_insts
 
 reifyFamilyInstance :: [Bool] -- True <=> the corresponding tv is poly-kinded
                               -- includes only *visible* tvs
@@ -1807,10 +1896,19 @@ reifyFamilyInstance is_poly_tvs (FamInst { fi_flavor = flavor
            ; th_tys <- reifyTypes types_only
            ; annot_th_tys <- zipWith3M annotThType is_poly_tvs types_only th_tys
            ; let lhs_type = mkThAppTs (TH.ConT fam') annot_th_tys
+           ; mb_sig <-
+               -- See "Wrinkle: data family instances with return kinds" in
+               -- Note [Reified instances and explicit kind signatures]
+               if (null cons || isGadtSyntaxTyCon rep_tc)
+                     && tyConAppNeedsKindSig False fam_tc (length ee_lhs)
+               then do { let full_kind = tcTypeKind (mkTyConApp fam_tc ee_lhs)
+                       ; th_full_kind <- reifyKind full_kind
+                       ; pure $ Just th_full_kind }
+               else pure Nothing
            ; return $
                if isNewTyCon rep_tc
-               then TH.NewtypeInstD [] th_tvs lhs_type Nothing (head cons) []
-               else TH.DataInstD    [] th_tvs lhs_type Nothing       cons  []
+               then TH.NewtypeInstD [] th_tvs lhs_type mb_sig (head cons) []
+               else TH.DataInstD    [] th_tvs lhs_type mb_sig       cons  []
            }
 
 ------------------------------
diff --git a/compiler/types/Type.hs b/compiler/types/Type.hs
index f8cd1da078..a5a93b0739 100644
--- a/compiler/types/Type.hs
+++ b/compiler/types/Type.hs
@@ -1775,7 +1775,6 @@ tyCoBinderVar_maybe (Named tv) = Just $ binderVar tv
 tyCoBinderVar_maybe _          = Nothing
 
 tyCoBinderType :: TyCoBinder -> Type
--- Barely used
 tyCoBinderType (Named tvb) = binderType tvb
 tyCoBinderType (Anon _ ty) = ty
 
diff --git a/testsuite/tests/th/T17296.hs b/testsuite/tests/th/T17296.hs
new file mode 100644
index 0000000000..a493923db1
--- /dev/null
+++ b/testsuite/tests/th/T17296.hs
@@ -0,0 +1,37 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeFamilies #-}
+module T17296 where
+
+import Data.Foldable
+import Data.Kind
+import Language.Haskell.TH hiding (Type)
+import System.IO
+
+data family   Foo1 :: Type -> Type
+data instance Foo1 Bool = Foo1Bool
+data instance Foo1 (Maybe a)
+
+data family   Foo2 :: k -> Type
+data instance Foo2 Bool = Foo2Bool
+data instance Foo2 (Maybe a)
+data instance Foo2 :: Char -> Type
+data instance Foo2 :: (Char -> Char) -> Type where
+
+data family   Foo3 :: k
+data instance Foo3
+data instance Foo3 Bool = Foo3Bool
+data instance Foo3 (Maybe a)
+data instance Foo3 :: Char -> Type
+data instance Foo3 :: (Char -> Char) -> Type where
+
+$(do let test :: Name -> Q ()
+         test n = do i <- reify n
+                     runIO $ do hPutStrLn stderr $ pprint i
+                                hPutStrLn stderr ""
+                                hFlush stderr
+
+     traverse_ test [''Foo1, ''Foo2, ''Foo3]
+     pure [])
diff --git a/testsuite/tests/th/T17296.stderr b/testsuite/tests/th/T17296.stderr
new file mode 100644
index 0000000000..4a6f1ac3bd
--- /dev/null
+++ b/testsuite/tests/th/T17296.stderr
@@ -0,0 +1,19 @@
+data family T17296.Foo1 :: * -> *
+data instance T17296.Foo1 GHC.Types.Bool = T17296.Foo1Bool
+data instance forall (a_0 :: *). T17296.Foo1 (GHC.Maybe.Maybe a_0)
+
+data family T17296.Foo2 :: k_0 -> *
+data instance T17296.Foo2 GHC.Types.Bool = T17296.Foo2Bool
+data instance forall (a_1 :: *). T17296.Foo2 (GHC.Maybe.Maybe a_1 :: *)
+data instance T17296.Foo2 :: GHC.Types.Char -> *
+data instance T17296.Foo2 :: (GHC.Types.Char -> GHC.Types.Char) ->
+                             *
+
+data family T17296.Foo3 :: k_0
+data instance T17296.Foo3 :: *
+data instance T17296.Foo3 GHC.Types.Bool = T17296.Foo3Bool
+data instance forall (a_1 :: *). T17296.Foo3 (GHC.Maybe.Maybe a_1 :: *)
+data instance T17296.Foo3 :: GHC.Types.Char -> *
+data instance T17296.Foo3 :: (GHC.Types.Char -> GHC.Types.Char) ->
+                             *
+
diff --git a/testsuite/tests/th/all.T b/testsuite/tests/th/all.T
index 590b060b0b..b8c0434595 100644
--- a/testsuite/tests/th/all.T
+++ b/testsuite/tests/th/all.T
@@ -485,3 +485,4 @@ test('T16976f', normal, compile_fail, [''])
 test('T16976z', normal, compile_fail, [''])
 test('T16980', normal, compile, [''])
 test('T16980a', normal, compile_fail, [''])
+test('T17296', normal, compile, ['-v0'])