Fix #16518 with some more kind-splitting smarts

This patch corrects two simple oversights that led to #16518:

1. `HsUtils.typeToLHsType` was taking visibility into account in the
   `TyConApp` case, but not the `AppTy` case. I've factored out the
   visibility-related logic into its own `go_app` function and now
   invoke `go_app` from both the `TyConApp` and `AppTy` cases.
2. `Type.fun_kind_arg_flags` did not properly split kinds with
   nested `forall`s, such as
   `(forall k. k -> Type) -> (forall k. k -> Type)`. This was simply
   because `fun_kind_arg_flags`'s `FunTy` case always bailed out and
   assumed all subsequent arguments were `Required`, which clearly
   isn't the case for nested `forall`s. I tweaked the `FunTy` case
   to recur on the result kind.

2 files changed, 38 insertions, 16 deletions

diff --git a/compiler/hsSyn/HsUtils.hs b/compiler/hsSyn/HsUtils.hs
index ea3c6aa197..531ff46ee4 100644
--- a/compiler/hsSyn/HsUtils.hs
+++ b/compiler/hsSyn/HsUtils.hs
@@ -106,7 +106,7 @@ import TcEvidence
 import RdrName
 import Var
 import TyCoRep
-import Type   ( tyConArgFlags )
+import Type   ( appTyArgFlags, splitAppTys, tyConArgFlags )
 import TysWiredIn ( unitTy )
 import TcType
 import DataCon
@@ -665,7 +665,6 @@ typeToLHsType ty
                           , hst_xforall = noExt
                           , hst_body = go tau })
     go (TyVarTy tv)         = nlHsTyVar (getRdrName tv)
-    go (AppTy t1 t2)        = nlHsAppTy (go t1) (go t2)
     go (LitTy (NumTyLit n))
       = noLoc $ HsTyLit NoExt (HsNumTy NoSourceText n)
     go (LitTy (StrTyLit s))
@@ -674,27 +673,35 @@ typeToLHsType ty
       | tyConAppNeedsKindSig True tc (length args)
         -- We must produce an explicit kind signature here to make certain
         -- programs kind-check. See Note [Kind signatures in typeToLHsType].
-      = nlHsParTy $ noLoc $ HsKindSig NoExt lhs_ty (go (tcTypeKind ty))
-      | otherwise = lhs_ty
+      = nlHsParTy $ noLoc $ HsKindSig NoExt ty' (go (tcTypeKind ty))
+      | otherwise = ty'
        where
-        arg_flags :: [ArgFlag]
-        arg_flags = tyConArgFlags tc args
-
-        lhs_ty :: LHsType GhcPs
-        lhs_ty = foldl' (\f (arg, flag) ->
-                          let arg' = go arg in
-                          case flag of
-                            Inferred  -> f
-                            Specified -> f `nlHsAppKindTy` arg'
-                            Required  -> f `nlHsAppTy`     arg')
-                        (nlHsTyVar (getRdrName tc))
-                        (zip args arg_flags)
+        ty' :: LHsType GhcPs
+        ty' = go_app (nlHsTyVar (getRdrName tc)) args (tyConArgFlags tc args)
+    go ty@(AppTy {})        = go_app (go head) args (appTyArgFlags head args)
+      where
+        head :: Type
+        args :: [Type]
+        (head, args) = splitAppTys ty
     go (CastTy ty _)        = go ty
     go (CoercionTy co)      = pprPanic "toLHsSigWcType" (ppr co)
 
          -- Source-language types have _invisible_ kind arguments,
          -- so we must remove them here (#8563)
 
+    go_app :: LHsType GhcPs -- The type being applied
+           -> [Type]        -- The argument types
+           -> [ArgFlag]     -- The argument types' visibilities
+           -> LHsType GhcPs
+    go_app head args arg_flags =
+      foldl' (\f (arg, flag) ->
+               let arg' = go arg in
+               case flag of
+                 Inferred  -> f
+                 Specified -> f `nlHsAppKindTy` arg'
+                 Required  -> f `nlHsAppTy`     arg')
+             head (zip args arg_flags)
+
     go_tv :: TyVar -> LHsTyVarBndr GhcPs
     go_tv tv = noLoc $ KindedTyVar noExt (noLoc (getRdrName tv))
                                    (go (tyVarKind tv))
diff --git a/compiler/types/Type.hs b/compiler/types/Type.hs
index 4426148967..c144d3e8f1 100644
--- a/compiler/types/Type.hs
+++ b/compiler/types/Type.hs
@@ -1698,6 +1698,21 @@ fun_kind_arg_flags = go emptyTCvSubst
         subst' = extendTvSubst subst tv arg_ty
     go subst (TyVarTy tv) arg_tys
       | Just ki <- lookupTyVar subst tv = go subst ki arg_tys
+    -- This FunTy case is important to handle kinds with nested foralls, such
+    -- as this kind (inspired by #16518):
+    --
+    --   forall {k1} k2. k1 -> k2 -> forall k3. k3 -> Type
+    --
+    -- Here, we want to get the following ArgFlags:
+    --
+    -- [Inferred,   Specified, Required, Required, Specified, Required]
+    -- forall {k1}. forall k2. k1 ->     k2 ->     forall k3. k3 ->     Type
+    go subst (FunTy{ft_af = af, ft_res = res_ki}) (_:arg_tys)
+      = argf : go subst res_ki arg_tys
+      where
+        argf = case af of
+                 VisArg   -> Required
+                 InvisArg -> Inferred
     go _ _ arg_tys = map (const Required) arg_tys
                         -- something is ill-kinded. But this can happen
                         -- when printing errors. Assume everything is Required.