WIP: make a smart mkAppTyMwip/T15952-2

This branch, wip/T15952-2, is WIP on the idea of making
a monadic mkAppTyM that ensures the Purely Kinded Invariant.

Needs comments etc.  But it validates all but one test!

9 files changed, 119 insertions, 28 deletions

diff --git a/compiler/typecheck/TcHsSyn.hs b/compiler/typecheck/TcHsSyn.hs
index 462b9245e5..3f7a32565d 100644
--- a/compiler/typecheck/TcHsSyn.hs
+++ b/compiler/typecheck/TcHsSyn.hs
@@ -204,6 +204,7 @@ data ZonkEnv  -- See Note [The ZonkEnv]
             , ze_tv_env :: TyCoVarEnv TyCoVar
             , ze_id_env :: IdEnv      Id
             , ze_meta_tv_env :: TcRef (TyVarEnv Type) }
+
 {- Note [The ZonkEnv]
 ~~~~~~~~~~~~~~~~~~~~~
 * ze_flexi :: ZonkFlexi says what to do with a
diff --git a/compiler/typecheck/TcHsType.hs b/compiler/typecheck/TcHsType.hs
index 006a97bd55..697fea750b 100644
--- a/compiler/typecheck/TcHsType.hs
+++ b/compiler/typecheck/TcHsType.hs
@@ -76,6 +76,7 @@ import TcUnify
 import TcIface
 import TcSimplify
 import TcHsSyn
+import TyCoRep  ( Type(..) )  -- Just for mkAppTyM
 import TcErrors ( reportAllUnsolved )
 import TcType
 import Inst   ( tcInstTyBinders, tcInstTyBinder )
@@ -970,7 +971,7 @@ tcInferApps mode orig_hs_ty fun_ty fun_ki orig_hs_args
   = do { traceTc "tcInferApps {" (ppr orig_hs_ty $$ ppr orig_hs_args $$ ppr fun_ki)
        ; (f_args, res_k) <- go 1 empty_subst fun_ty orig_ki_binders orig_inner_ki orig_hs_args
        ; traceTc "tcInferApps }" empty
-       ; res_k <- zonkTcType res_k  -- Uphold (IT4) of Note [The tcType invariant]
+--       ; res_k <- zonkTcType res_k  -- Uphold (IT4) of Note [The tcType invariant]
        ; return (f_args, res_k) }
   where
     empty_subst                      = mkEmptyTCvSubst $ mkInScopeSet $
@@ -1010,16 +1011,16 @@ tcInferApps mode orig_hs_ty fun_ty fun_ki orig_hs_args
                             tc_lhs_type (kindLevel mode) ki exp_kind
                   ; traceTc "tcInferApps (vis kind app)" (ppr exp_kind)
                   ; let subst' = extendTvSubstBinderAndInScope subst ki_binder arg'
-                  ; go (n+1) subst'
-                       (mkNakedAppTy fun arg')
+                  ; fun' <- mkAppTyM fun arg'
+                  ; go (n+1) subst' fun'
                        ki_binders inner_ki args }
 
       | isInvisibleBinder ki_binder
           -- Instantiate if not specified or if there is no kind application
       = do { traceTc "tcInferApps (invis normal app)" (ppr ki_binder $$ ppr subst $$ ppr (tyCoBinderArgFlag ki_binder))
            ; (subst', arg') <- tcInstTyBinder Nothing subst ki_binder
-           ; go n subst' (mkNakedAppTy fun arg')
-                        ki_binders inner_ki all_args }
+           ; fun' <- mkAppTyM fun arg'
+           ; go n subst' fun' ki_binders inner_ki all_args }
 
       | otherwise -- if binder is visible
          = case arg of
@@ -1035,9 +1036,9 @@ tcInferApps mode orig_hs_ty fun_ty fun_ki orig_hs_args
                                tc_lhs_type mode ty exp_kind
                      ; traceTc "tcInferApps (vis normal app)" (ppr exp_kind)
                      ; let subst' = extendTvSubstBinderAndInScope subst ki_binder arg'
-                     ; go (n+1) subst'
-                          (mkNakedAppTy fun arg')
-                          ki_binders inner_ki args }
+                     ; fun' <- mkAppTyM fun arg'
+                     ; go (n+1) subst' fun' ki_binders inner_ki args }
+
             -- error if the argument is a kind application
              HsTypeArg ki -> do { traceTc "tcInferApps (error)"
                                     (vcat [ ppr ki_binder
@@ -1079,11 +1080,52 @@ tcInferApps mode orig_hs_ty fun_ty fun_ki orig_hs_args
     ty_app_err arg ty = failWith $ text "Cannot apply function of kind" <+> quotes (ppr ty)
                            $$ text "to visible kind argument" <+> quotes (ppr arg)
 
+mkAppTyM :: TcType -> TcType -> TcM TcType
+-- If 'fun' and 'arg' satisfy the Purely Kinded Invariant,
+-- and, if fully-zonked (fun arg) is well-kinded, then
+-- the result satisfies the Purely Kinded Invariant
+-- See Note [mkAppTyM]
+mkAppTyM fun arg
+  | TyConApp tc args <- fun
+  , isTypeSynonymTyCon tc
+  , args `lengthIs` (tyConArity tc - 1)
+  = do { traceTc "mkAppTyM synonym" (ppr fun <+> dcolon <+> ppr (tcTypeKind fun))
+       ; arg'  <- zonkTcType  arg
+       ; args' <- zonkTcTypes args
+       ; return (TyConApp tc (args' ++ [arg'])) }
+
+  | isPiTy (tcTypeKind fun)
+  = do { traceTc "mkAppTyM no zonk 1" (ppr fun)
+       ; traceTc "mkAppTyM no zonk 2" (ppr fun <+> dcolon <+> ppr (tcTypeKind fun))
+       ; traceTc "mkAppTyM no zonk" (vcat [ ppr fun <+> dcolon <+> ppr (tcTypeKind fun)
+                                          , ppr (mkAppTy fun arg) ])
+       ; return (mkAppTy fun arg) }
+
+  | otherwise
+  = do { fun' <- zonkTcType fun
+       ; traceTc "mkAppTyM with zonk" (vcat [ ppr fun' <+> dcolon <+> ppr (tcTypeKind fun')
+                                            , ppr (mkAppTy fun' arg) ])
+       ; return (mkAppTy fun' arg) }
+
+{- Note [mkAppTyM]
+~~~~~~~~~~~~~~~~~~
+
+Nasty case
+  type S f a = f a
+
+The type (S ff aa) may not satisfy the Purely Kinded Invariant because its expansion
+is (ff aa), and that does not satisfy PKI.  E.g. perhaps (ff :: kappa), where
+'kappa' has already been unified with (*->*).
+
+Solution: when building a saturated type synonym, fully zonk it.
+-}
+
+
 appTypeToArg :: LHsType GhcRn -> [LHsTypeArg GhcRn] -> LHsType GhcRn
-appTypeToArg f [] = f
-appTypeToArg f (HsValArg arg : args) = appTypeToArg (mkHsAppTy f arg) args
+appTypeToArg f []                     = f
+appTypeToArg f (HsValArg arg  : args) = appTypeToArg (mkHsAppTy f arg)     args
 appTypeToArg f (HsTypeArg arg : args) = appTypeToArg (mkHsAppKindTy f arg) args
-appTypeToArg f (HsArgPar _ : arg) = appTypeToArg f arg
+appTypeToArg f (HsArgPar _    : args) = appTypeToArg f                     args
 
 -- | Applies a type to a list of arguments.
 -- Always consumes all the arguments, using 'matchExpectedFunKind' as
@@ -1131,7 +1173,8 @@ checkExpectedKind :: HasDebugCallStack
                   -> TcKind         -- ^ the expected kind
                   -> TcM TcType
 checkExpectedKind sat hs_ty ty act exp
-  = do { (new_ty, new_act) <- case splitTyConApp_maybe ty of
+  = do { traceTc "checkExpectedKind" (ppr ty $$ ppr act)
+       ; (new_ty, new_act) <- case splitTyConApp_maybe ty of
            Just (tc, args)
              -- if the family tycon must be saturated and is not yet satured
              -- If we don't do this, we get #11246
diff --git a/compiler/typecheck/TcMType.hs b/compiler/typecheck/TcMType.hs
index ffeb602382..35d2b4578a 100644
--- a/compiler/typecheck/TcMType.hs
+++ b/compiler/typecheck/TcMType.hs
@@ -1937,7 +1937,7 @@ zonkTcTypeMapper = TyCoMapper
   , tcm_covar = const (\cv -> mkCoVarCo <$> zonkTyCoVarKind cv)
   , tcm_hole  = hole
   , tcm_tycobinder = \_env tv _vis -> ((), ) <$> zonkTyCoVarKind tv
-  , tcm_tycon = return }
+  , tcm_tycon = zonk_tc_tycon }
   where
     hole :: () -> CoercionHole -> TcM Coercion
     hole _ hole@(CoercionHole { ch_ref = ref, ch_co_var = cv })
@@ -1948,6 +1948,9 @@ zonkTcTypeMapper = TyCoMapper
                Nothing -> do { cv' <- zonkCoVar cv
                              ; return $ HoleCo (hole { ch_co_var = cv' }) } }
 
+    zonk_tc_tycon tc = do { tck' <- zonkTcType (tyConKind tc)
+                          ; return (setTcTyConKind tc tck') }
+    
 -- For unbound, mutable tyvars, zonkType uses the function given to it
 -- For tyvars bound at a for-all, zonkType zonks them to an immutable
 --      type variable and zonks the kind too
diff --git a/compiler/typecheck/TcType.hs b/compiler/typecheck/TcType.hs
index b2c9b3291f..0a264a30bc 100644
--- a/compiler/typecheck/TcType.hs
+++ b/compiler/typecheck/TcType.hs
@@ -225,7 +225,7 @@ import ErrUtils( Validity(..), MsgDoc, isValid )
 import qualified GHC.LanguageExtensions as LangExt
 
 import Data.List  ( mapAccumL )
-import Data.Functor.Identity( Identity(..) )
+-- import Data.Functor.Identity( Identity(..) )
 import Data.IORef
 import Data.List.NonEmpty( NonEmpty(..) )
 
@@ -1351,6 +1351,21 @@ Notes:
 
 ---------------
 mkNakedAppTys :: Type -> [Type] -> Type
+mkNakedAppTys = mkAppTys
+
+mkNakedAppTy :: Type -> Type -> Type
+mkNakedAppTy = mkAppTy
+
+mkNakedCastTy :: Type -> Coercion -> Type
+mkNakedCastTy = mkCastTy
+ -- Do we need a tc version of mkCastTy?
+
+nakedSubstTy :: HasCallStack => TCvSubst -> TcType  -> TcType
+nakedSubstTy = substTy
+
+
+{-
+mkNakedAppTys :: Type -> [Type] -> Type
 -- See Note [The well-kinded type invariant]
 mkNakedAppTys ty1                []   = ty1
 mkNakedAppTys (TyConApp tc tys1) tys2 = mkTyConApp tc (tys1 ++ tys2)
@@ -1387,6 +1402,7 @@ nakedSubstMapper
                , tcm_hole       = \_ hole   -> return (HoleCo hole)
                , tcm_tycobinder = \subst tv _ -> return (substVarBndr subst tv)
                , tcm_tycon    = return }
+-}
 
 {-
 ************************************************************************
diff --git a/compiler/types/TyCon.hs b/compiler/types/TyCon.hs
index eb0b84d47e..17cc8e64f6 100644
--- a/compiler/types/TyCon.hs
+++ b/compiler/types/TyCon.hs
@@ -69,7 +69,8 @@ module TyCon(
         isTyConAssoc, tyConAssoc_maybe, tyConFlavourAssoc_maybe,
         isImplicitTyCon,
         isTyConWithSrcDataCons,
-        isTcTyCon, isTcLevPoly,
+        isTcTyCon, setTcTyConKind,
+        isTcLevPoly,
 
         -- ** Extracting information out of TyCons
         tyConName,
@@ -2131,6 +2132,14 @@ isTcTyCon :: TyCon -> Bool
 isTcTyCon (TcTyCon {}) = True
 isTcTyCon _            = False
 
+setTcTyConKind :: TyCon -> Kind -> TyCon
+-- Update the Kind of a TcTyCon
+-- The new kind is always a zonked version of its previous
+-- kind, so we don't need to update any other fields.
+-- See Note [The Purely Kinded Invariant] in TcHsType
+setTcTyConKind tc@(TcTyCon {}) kind = tc { tyConKind = kind }
+setTcTyConKind tc              _    = pprPanic "setTcTyConKind" (ppr tc)
+
 -- | Could this TyCon ever be levity-polymorphic when fully applied?
 -- True is safe. False means we're sure. Does only a quick check
 -- based on the TyCon's category.
diff --git a/compiler/types/Type.hs b/compiler/types/Type.hs
index 2fd060a7f4..537e5808fe 100644
--- a/compiler/types/Type.hs
+++ b/compiler/types/Type.hs
@@ -538,9 +538,11 @@ data TyCoMapper env m
           -- ^ The returned env is used in the extended scope
 
       , tcm_tycon :: TyCon -> m TyCon
-          -- ^ This is used only to turn 'TcTyCon's into 'TyCon's.
-          -- See Note [Type checking recursive type and class declarations]
-          -- in TcTyClsDecls
+          -- ^ This is used only for TcTyCons
+          -- a) To zonk TcTyCons
+          -- b) To turn TcTyCons into TyCons.
+          --    See Note [Type checking recursive type and class declarations]
+          --    in TcTyClsDecls
       }
 
 {-# INLINABLE mapType #-}  -- See Note [Specialising mappers]
@@ -552,12 +554,18 @@ mapType mapper@(TyCoMapper { tcm_smart = smart, tcm_tyvar = tyvar
   where
     go (TyVarTy tv) = tyvar env tv
     go (AppTy t1 t2) = mkappty <$> go t1 <*> go t2
-    go t@(TyConApp tc []) | not (isTcTyCon tc)
-                          = return t  -- avoid allocation in this exceedingly
-                                      -- common case (mostly, for *)
-    go (TyConApp tc tys)
+    go ty@(TyConApp tc tys)
+      | isTcTyCon tc
       = do { tc' <- tycon tc
            ; mktyconapp tc' <$> mapM go tys }
+
+      -- Not a TcTyCon
+      | null tys    -- Avoid allocation in this very
+      = return ty   -- common case (E.g. Int, LiftedRep etc)
+
+      | otherwise
+      = mktyconapp tc <$> mapM go tys 
+
     go (FunTy arg res)   = FunTy <$> go arg <*> go res
     go (ForAllTy (Bndr tv vis) inner)
       = do { (env', tv') <- tycobinder env tv vis
@@ -586,7 +594,9 @@ mapCoercion mapper@(TyCoMapper { tcm_smart = smart, tcm_covar = covar
     go (Refl ty) = Refl <$> mapType mapper env ty
     go (GRefl r ty mco) = mkgreflco r <$> mapType mapper env ty <*> (go_mco mco)
     go (TyConAppCo r tc args)
-      = do { tc' <- tycon tc
+      = do { tc' <- if isTcTyCon tc
+                    then tycon tc
+                    else return tc
            ; mktyconappco r tc' <$> mapM go args }
     go (AppCo c1 c2) = mkappco <$> go c1 <*> go c2
     go (ForAllCo tv kind_co co)
@@ -719,6 +729,10 @@ mkAppTy ty1               ty2 = AppTy ty1 ty2
         --
         -- Here Id is partially applied in the type sig for Foo,
         -- but once the type synonyms are expanded all is well
+        --
+        -- Moreover in TcHsTypes.tcInferApps we build up a type
+        --   (T t1 t2 t3) one argument at a type, thus forming
+        --   (T t1), (T t1 t2), etc
 
 mkAppTys :: Type -> [Type] -> Type
 mkAppTys ty1                []   = ty1
@@ -1010,6 +1024,8 @@ piResultTy ty arg = case piResultTy_maybe ty arg of
                       Nothing  -> pprPanic "piResultTy" (ppr ty $$ ppr arg)
 
 piResultTy_maybe :: Type -> Type -> Maybe Type
+-- We don't need a 'tc' version, because
+-- this function behaves the same for Type and Constraint
 piResultTy_maybe ty arg
   | Just ty' <- coreView ty = piResultTy_maybe ty' arg
 
@@ -2763,7 +2779,7 @@ tcTypeKind ty@(ForAllTy {})
     body_kind = tcTypeKind body
 
 
-isPredTy :: Type -> Bool
+isPredTy :: HasDebugCallStack => Type -> Bool
 -- See Note [Types for coercions, predicates, and evidence]
 isPredTy ty = tcIsConstraintKind (tcTypeKind ty)
 
diff --git a/compiler/types/Type.hs-boot b/compiler/types/Type.hs-boot
index 84fbaca6f7..d9ea519a08 100644
--- a/compiler/types/Type.hs-boot
+++ b/compiler/types/Type.hs-boot
@@ -8,7 +8,7 @@ import Var ( TyCoVar )
 import {-# SOURCE #-} TyCoRep( Type, Coercion )
 import Util
 
-isPredTy     :: Type -> Bool
+isPredTy     :: HasDebugCallStack => Type -> Bool
 isCoercionTy :: Type -> Bool
 
 mkAppTy    :: Type -> Type -> Type
diff --git a/compiler/utils/Outputable.hs b/compiler/utils/Outputable.hs
index bb3b9d3177..68a189f3c0 100644
--- a/compiler/utils/Outputable.hs
+++ b/compiler/utils/Outputable.hs
@@ -1199,7 +1199,7 @@ pprTraceException heading doc =
 pprSTrace :: HasCallStack => SDoc -> a -> a
 pprSTrace doc = pprTrace "" (doc $$ callStackDoc)
 
-warnPprTrace :: Bool -> String -> Int -> SDoc -> a -> a
+warnPprTrace :: HasCallStack => Bool -> String -> Int -> SDoc -> a -> a
 -- ^ Just warn about an assertion failure, recording the given file and line number.
 -- Should typically be accessed with the WARN macros
 warnPprTrace _     _     _     _    x | not debugIsOn     = x
@@ -1207,7 +1207,9 @@ warnPprTrace _     _file _line _msg x
    | hasNoDebugOutput unsafeGlobalDynFlags = x
 warnPprTrace False _file _line _msg x = x
 warnPprTrace True   file  line  msg x
-  = pprDebugAndThen unsafeGlobalDynFlags trace heading msg x
+  = pprDebugAndThen unsafeGlobalDynFlags trace heading
+                    (msg $$ callStackDoc )
+                    x
   where
     heading = hsep [text "WARNING: file", text file <> comma, text "line", int line]
 
diff --git a/compiler/utils/Outputable.hs-boot b/compiler/utils/Outputable.hs-boot
index ad7d091833..fb3c173a33 100644
--- a/compiler/utils/Outputable.hs-boot
+++ b/compiler/utils/Outputable.hs-boot
@@ -1,11 +1,12 @@
 module Outputable where
 
 import GhcPrelude
+import GHC.Stack( HasCallStack )
 
 data SDoc
 
 showSDocUnsafe :: SDoc -> String
 
-warnPprTrace :: Bool -> String -> Int -> SDoc -> a -> a
+warnPprTrace :: HasCallStack => Bool -> String -> Int -> SDoc -> a -> a
 
 text :: String -> SDoc