Improve piResultTys and friends

Several things here:

* Re-implement piResultTys so that its substitution has
  the correct in-scope set

  That means paying close attention to performance, since as we
  discovered in Trac #11371, it's a heavily used function and
  is often used on ordinary function types, with no foralls to
  worry about substituting.

* Kill off applyTys, which was just the same as piResultTys.

* Re-engineer MkCore.mkCoreApps so that it calls piResultTys,
  rather than repeatedly calling piResultTy.

9 files changed, 110 insertions, 92 deletions

diff --git a/compiler/coreSyn/CoreLint.hs b/compiler/coreSyn/CoreLint.hs
index 1d4d28c151..6dbcbe4ce9 100644
--- a/compiler/coreSyn/CoreLint.hs
+++ b/compiler/coreSyn/CoreLint.hs
@@ -912,7 +912,7 @@ lintCoreAlt scrut_ty alt_ty alt@(DataAlt con, args, rhs)
         -- type variables of the data constructor
         -- We've already check
       lintL (tycon == dataConTyCon con) (mkBadConMsg tycon con)
-    ; let con_payload_ty = applyTys (dataConRepType con) tycon_arg_tys
+    ; let con_payload_ty = piResultTys (dataConRepType con) tycon_arg_tys
 
         -- And now bring the new binders into scope
     ; lintBinders args $ \ args' -> do
diff --git a/compiler/coreSyn/CoreUtils.hs b/compiler/coreSyn/CoreUtils.hs
index 6fa55c91a3..d2da53a97a 100644
--- a/compiler/coreSyn/CoreUtils.hs
+++ b/compiler/coreSyn/CoreUtils.hs
@@ -175,15 +175,15 @@ applyTypeToArgs e op_ty args
                                   = go res_ty args
     go _ _ = pprPanic "applyTypeToArgs" panic_msg
 
-    -- go_ty_args: accumulate type arguments so we can instantiate all at once
+    -- go_ty_args: accumulate type arguments so we can
+    -- instantiate all at once with piResultTys
     go_ty_args op_ty rev_tys (Type ty : args)
        = go_ty_args op_ty (ty:rev_tys) args
     go_ty_args op_ty rev_tys (Coercion co : args)
        = go_ty_args op_ty (mkCoercionTy co : rev_tys) args
     go_ty_args op_ty rev_tys args
-       = go (applyTysD panic_msg_w_hdr op_ty (reverse rev_tys)) args
+       = go (piResultTys op_ty (reverse rev_tys)) args
 
-    panic_msg_w_hdr = hang (text "applyTypeToArgs") 2 panic_msg
     panic_msg = vcat [ text "Expression:" <+> pprCoreExpr e
                      , text "Type:" <+> ppr op_ty
                      , text "Args:" <+> ppr args ]
diff --git a/compiler/coreSyn/MkCore.hs b/compiler/coreSyn/MkCore.hs
index e012f2c4be..94a264c120 100644
--- a/compiler/coreSyn/MkCore.hs
+++ b/compiler/coreSyn/MkCore.hs
@@ -120,13 +120,13 @@ mkCoreLets binds body = foldr mkCoreLet body binds
 mkCoreApp :: SDoc -> CoreExpr -> CoreExpr -> CoreExpr
 -- Respects the let/app invariant by building a case expression where necessary
 --   See CoreSyn Note [CoreSyn let/app invariant]
-mkCoreApp _ fun (Type ty) = App fun (Type ty)
+mkCoreApp _ fun (Type ty)     = App fun (Type ty)
 mkCoreApp _ fun (Coercion co) = App fun (Coercion co)
-mkCoreApp d fun arg       = ASSERT2( isFunTy fun_ty, ppr fun $$ ppr arg $$ d )
-                          mk_val_app fun arg arg_ty res_ty
-                      where
-                        fun_ty = exprType fun
-                        (arg_ty, res_ty) = splitFunTy fun_ty
+mkCoreApp d fun arg           = ASSERT2( isFunTy fun_ty, ppr fun $$ ppr arg $$ d )
+                                mk_val_app fun arg arg_ty res_ty
+                              where
+                                fun_ty = exprType fun
+                                (arg_ty, res_ty) = splitFunTy fun_ty
 
 -- | Construct an expression which represents the application of a number of
 -- expressions to another. The leftmost expression in the list is applied first
@@ -137,13 +137,13 @@ mkCoreApps :: CoreExpr -> [CoreExpr] -> CoreExpr
 mkCoreApps orig_fun orig_args
   = go orig_fun (exprType orig_fun) orig_args
   where
-    go fun _      []                   = fun
-    go fun fun_ty (Type ty     : args) = go (App fun (Type ty)) (piResultTy fun_ty ty) args
-    go fun fun_ty (arg         : args) = ASSERT2( isFunTy fun_ty, ppr fun_ty $$ ppr orig_fun
-                                                                  $$ ppr orig_args )
-                                         go (mk_val_app fun arg arg_ty res_ty) res_ty args
-                                       where
-                                         (arg_ty, res_ty) = splitFunTy fun_ty
+    go fun _      []               = fun
+    go fun fun_ty (Type ty : args) = go (App fun (Type ty)) (piResultTy fun_ty ty) args
+    go fun fun_ty (arg     : args) = ASSERT2( isFunTy fun_ty, ppr fun_ty $$ ppr orig_fun
+                                                              $$ ppr orig_args )
+                                     go (mk_val_app fun arg arg_ty res_ty) res_ty args
+                                   where
+                                     (arg_ty, res_ty) = splitFunTy fun_ty
 
 -- | Construct an expression which represents the application of a number of
 -- expressions to that of a data constructor expression. The leftmost expression
diff --git a/compiler/iface/BuildTyCl.hs b/compiler/iface/BuildTyCl.hs
index bde9019ce1..0022e29f11 100644
--- a/compiler/iface/BuildTyCl.hs
+++ b/compiler/iface/BuildTyCl.hs
@@ -74,7 +74,7 @@ mkNewTyConRhs tycon_name tycon con
   where
     tvs    = tyConTyVars tycon
     roles  = tyConRoles tycon
-    inst_con_ty = applyTys (dataConUserType con) (mkTyVarTys tvs)
+    inst_con_ty = piResultTys (dataConUserType con) (mkTyVarTys tvs)
     rhs_ty = ASSERT( isFunTy inst_con_ty ) funArgTy inst_con_ty
         -- Instantiate the data con with the
         -- type variables from the tycon
diff --git a/compiler/typecheck/TcClassDcl.hs b/compiler/typecheck/TcClassDcl.hs
index 8a67965234..b1baabb963 100644
--- a/compiler/typecheck/TcClassDcl.hs
+++ b/compiler/typecheck/TcClassDcl.hs
@@ -26,7 +26,7 @@ import TcBinds
 import TcUnify
 import TcHsType
 import TcMType
-import Type     ( getClassPredTys_maybe, varSetElemsWellScoped )
+import Type     ( getClassPredTys_maybe, varSetElemsWellScoped, piResultTys )
 import TcType
 import TcRnMonad
 import BuildTyCl( TcMethInfo )
@@ -299,7 +299,7 @@ instantiateMethod :: Class -> Id -> [TcType] -> TcType
 instantiateMethod clas sel_id inst_tys
   = ASSERT( ok_first_pred ) local_meth_ty
   where
-    rho_ty = applyTys (idType sel_id) inst_tys
+    rho_ty = piResultTys (idType sel_id) inst_tys
     (first_pred, local_meth_ty) = tcSplitPredFunTy_maybe rho_ty
                 `orElse` pprPanic "tcInstanceMethod" (ppr sel_id)
 
diff --git a/compiler/typecheck/TcInstDcls.hs b/compiler/typecheck/TcInstDcls.hs
index fdc9e8d00e..82c66cc953 100644
--- a/compiler/typecheck/TcInstDcls.hs
+++ b/compiler/typecheck/TcInstDcls.hs
@@ -1280,7 +1280,7 @@ tcMethods dfun_id clas tyvars dfun_ev_vars inst_tys
                               nO_METHOD_BINDING_ERROR_ID
         error_msg dflags = L inst_loc (HsLit (HsStringPrim ""
                                               (unsafeMkByteString (error_string dflags))))
-        meth_tau     = funResultTy (applyTys (idType sel_id) inst_tys)
+        meth_tau     = funResultTy (piResultTys (idType sel_id) inst_tys)
         error_string dflags = showSDoc dflags
                               (hcat [ppr inst_loc, vbar, ppr sel_id ])
         lam_wrapper  = mkWpTyLams tyvars <.> mkWpLams dfun_ev_vars
diff --git a/compiler/typecheck/TcType.hs b/compiler/typecheck/TcType.hs
index 0160310313..8021c75a22 100644
--- a/compiler/typecheck/TcType.hs
+++ b/compiler/typecheck/TcType.hs
@@ -131,7 +131,7 @@ module TcType (
 
   mkForAllTy, mkForAllTys, mkInvForAllTys, mkSpecForAllTys, mkNamedForAllTy,
   mkFunTy, mkFunTys,
-  mkTyConApp, mkAppTy, mkAppTys, applyTys,
+  mkTyConApp, mkAppTy, mkAppTys,
   mkTyConTy, mkTyVarTy,
   mkTyVarTys,
   mkNamedBinder,
diff --git a/compiler/types/Coercion.hs b/compiler/types/Coercion.hs
index 3d841e5bea..7473fb9663 100644
--- a/compiler/types/Coercion.hs
+++ b/compiler/types/Coercion.hs
@@ -1766,7 +1766,7 @@ coercionKind co = go co
     -- Collect up all the arguments and apply all at once
     -- See Note [Nested InstCos]
     go_app (InstCo co arg) args = go_app co (arg:args)
-    go_app co              args = applyTys <$> go co <*> (sequenceA $ map go args)
+    go_app co              args = piResultTys <$> go co <*> (sequenceA $ map go args)
 
     -- The real mkCastTy is too slow, and we can easily have nested ForAllCos.
     mk_cast_ty :: Type -> Coercion -> Type
@@ -1831,7 +1831,7 @@ coercionKindRole = go
     go_app (InstCo co arg) args = go_app co (arg:args)
     go_app co              args
       = let (pair, r) = go co in
-        (applyTys <$> pair <*> (sequenceA $ map coercionKind args), r)
+        (piResultTys <$> pair <*> (sequenceA $ map coercionKind args), r)
 
 -- | Retrieve the role from a coercion.
 coercionRole :: Coercion -> Role
@@ -1852,8 +1852,7 @@ If we deal with the InstCos one at a time, we'll do this:
    1.  Find the kind of (g @ ty1 .. @ ty99) : forall a100. phi'
    2.  Substitute phi'[ ty100/a100 ], a single tyvar->type subst
 But this is a *quadratic* algorithm, and the blew up Trac #5631.
-So it's very important to do the substitution simultaneously.
-
-cf Type.applyTys (which in fact we call here)
+So it's very important to do the substitution simultaneously;
+cf Type.piResultTys (which in fact we call here).
 
 -}
diff --git a/compiler/types/Type.hs b/compiler/types/Type.hs
index c8f41db699..824aa9d752 100644
--- a/compiler/types/Type.hs
+++ b/compiler/types/Type.hs
@@ -43,7 +43,7 @@ module Type (
         splitNamedPiTys,
         mkPiType, mkPiTypes, mkPiTypesPreferFunTy,
         piResultTy, piResultTys,
-        applyTys, applyTysD, applyTysX, dropForAlls,
+        applyTysX, dropForAlls,
 
         mkNumLitTy, isNumLitTy,
         mkStrLitTy, isStrLitTy,
@@ -812,19 +812,9 @@ splitFunTysN n ty = ASSERT2( isFunTy ty, int n <+> ppr ty )
 
 funResultTy :: Type -> Type
 -- ^ Extract the function result type and panic if that is not possible
-funResultTy ty = piResultTy ty (pprPanic "funResultTy" (ppr ty))
-
--- | Essentially 'funResultTy' on kinds handling pi-types too
-piResultTy :: Type -> Type -> Type
-piResultTy ty arg | Just ty' <- coreView ty = piResultTy ty' arg
-piResultTy (ForAllTy (Anon _) res)     _   = res
-piResultTy (ForAllTy (Named tv _) res) arg = substTyWithUnchecked [tv] [arg] res
-piResultTy ty arg                          = pprPanic "piResultTy"
-                                                 (ppr ty $$ ppr arg)
-
--- | Fold 'piResultTy' over many types
-piResultTys :: Type -> [Type] -> Type
-piResultTys = foldl piResultTy
+funResultTy ty | Just ty' <- coreView ty = funResultTy ty'
+funResultTy (ForAllTy (Anon {}) res)     = res
+funResultTy ty                           = pprPanic "funResultTy" (ppr ty)
 
 funArgTy :: Type -> Type
 -- ^ Extract the function argument type and panic if that is not possible
@@ -832,6 +822,80 @@ funArgTy ty | Just ty' <- coreView ty = funArgTy ty'
 funArgTy (ForAllTy (Anon arg) _res) = arg
 funArgTy ty                         = pprPanic "funArgTy" (ppr ty)
 
+piResultTy :: Type -> Type -> Type
+-- ^ Just like 'piResultTys' but for a single argument
+-- Try not to iterate 'piResultTy', because it's inefficient to substitute
+-- one variable at a time; instead use 'piResultTys"
+piResultTy ty arg
+  | Just ty' <- coreView ty = piResultTy ty' arg
+
+  | ForAllTy bndr res <- ty
+  = case bndr of
+      Anon {}    -> res
+      Named tv _ -> substTy (extendTvSubst empty_subst tv arg) res
+        where
+          empty_subst = mkEmptyTCvSubst $ mkInScopeSet $
+                        tyCoVarsOfTypes [arg,res]
+  | otherwise
+  = panic "piResultTys"
+
+-- | (piResultTys f_ty [ty1, .., tyn]) gives the type of (f ty1 .. tyn)
+--   where f :: f_ty
+-- 'piResultTys' is interesting because:
+--      1. 'f_ty' may have more for-alls than there are args
+--      2. Less obviously, it may have fewer for-alls
+-- For case 2. think of:
+--   piResultTys (forall a.a) [forall b.b, Int]
+-- This really can happen, but only (I think) in situations involving
+-- undefined.  For example:
+--       undefined :: forall a. a
+-- Term: undefined @(forall b. b->b) @Int
+-- This term should have type (Int -> Int), but notice that
+-- there are more type args than foralls in 'undefined's type.
+
+-- If you edit this function, you may need to update the GHC formalism
+-- See Note [GHC Formalism] in coreSyn/CoreLint.hs
+
+-- This is a heavily used function (e.g. from typeKind),
+-- so we pay attention to efficiency, especially in the special case
+-- where there are no for-alls so we are just dropping arrows from
+-- a function type/kind.
+piResultTys :: Type -> [Type] -> Type
+piResultTys ty [] = ty
+piResultTys ty orig_args@(arg:args)
+  | Just ty' <- coreView ty
+  = piResultTys ty' orig_args
+
+  | ForAllTy bndr res <- ty
+  = case bndr of
+      Anon {}    -> piResultTys res args
+      Named tv _ -> go (extendVarEnv emptyTvSubstEnv tv arg) res args
+
+  | otherwise
+  = panic "piResultTys"
+  where
+    go :: TvSubstEnv -> Type -> [Type] -> Type
+    go tv_env ty [] = substTy (mkTvSubst in_scope tv_env) ty
+      where
+        in_scope = mkInScopeSet (tyCoVarsOfTypes (ty:orig_args))
+
+    go tv_env ty all_args@(arg:args)
+      | Just ty' <- coreView ty
+      = go tv_env ty' all_args
+
+      | ForAllTy bndr res <- ty
+      = case bndr of
+          Anon _     -> go tv_env res args
+          Named tv _ -> go (extendVarEnv tv_env tv arg) res args
+
+      | TyVarTy tv <- ty
+      , Just ty' <- lookupVarEnv tv_env tv
+        -- Deals with piResultTys (forall a. a) [forall b.b, Int]
+      = piResultTys ty' all_args
+
+      | otherwise
+      = panic "piResultTys"
+
 {-
 ---------------------------------------------------------------------
                                 TyConApp
@@ -1393,61 +1457,16 @@ splitPiTysInvisible ty = split ty ty []
 tyConBinders :: TyCon -> [TyBinder]
 tyConBinders = fst . splitPiTys . tyConKind
 
-{-
-applyTys
-~~~~~~~~~~~~~~~~~
--}
-
-applyTys :: Type -> [KindOrType] -> Type
--- ^ This function is interesting because:
---
---      1. The function may have more for-alls than there are args
---
---      2. Less obviously, it may have fewer for-alls
---
--- For case 2. think of:
---
--- > applyTys (forall a.a) [forall b.b, Int]
---
--- This really can happen, but only (I think) in situations involving
--- undefined.  For example:
---       undefined :: forall a. a
--- Term: undefined @(forall b. b->b) @Int
--- This term should have type (Int -> Int), but notice that
--- there are more type args than foralls in 'undefined's type.
-
--- If you edit this function, you may need to update the GHC formalism
--- See Note [GHC Formalism] in coreSyn/CoreLint.hs
-applyTys ty args = applyTysD empty ty args
-
-applyTysD :: SDoc -> Type -> [Type] -> Type     -- Debug version
-applyTysD doc orig_fun_ty arg_tys
-  | null arg_tys
-  = orig_fun_ty
-
-  | n_bndrs == n_args     -- The vastly common case
-  = substTyWithBindersUnchecked bndrs arg_tys rho_ty
-
-  | n_bndrs > n_args      -- Too many for-alls
-  = substTyWithBindersUnchecked (take n_args bndrs) arg_tys
-                                (mkForAllTys (drop n_args bndrs) rho_ty)
-
-  | otherwise             -- Too many type args
-  = ASSERT2( n_bndrs > 0, doc $$ ppr orig_fun_ty $$ ppr arg_tys )       -- Zero case gives infinite loop!
-    applyTysD doc (substTyWithBinders bndrs (take n_bndrs arg_tys) rho_ty)
-                  (drop n_bndrs arg_tys)
-  where
-    (bndrs, rho_ty) = splitPiTys orig_fun_ty
-    n_bndrs = length bndrs
-    n_args  = length arg_tys
-
 applyTysX :: [TyVar] -> Type -> [Type] -> Type
 -- applyTyxX beta-reduces (/\tvs. body_ty) arg_tys
+-- Assumes that (/\tvs. body_ty) is closed
 applyTysX tvs body_ty arg_tys
-  = ASSERT2( length arg_tys >= n_tvs, ppr tvs $$ ppr body_ty $$ ppr arg_tys )
-    mkAppTys (substTyWithUnchecked tvs (take n_tvs arg_tys) body_ty)
+  = ASSERT2( length arg_tys >= n_tvs, pp_stuff )
+    ASSERT2( tyCoVarsOfType body_ty `subVarSet` mkVarSet tvs, pp_stuff )
+    mkAppTys (substTyWith tvs (take n_tvs arg_tys) body_ty)
              (drop n_tvs arg_tys)
   where
+    pp_stuff = vcat [ppr tvs, ppr body_ty, ppr arg_tys]
     n_tvs = length tvs
 
 {-
@@ -1578,7 +1597,7 @@ isPredTy ty = go ty []
 
     go_k :: Kind -> [KindOrType] -> Bool
     -- True <=> ('k' applied to 'kts') = Constraint
-    go_k k args = isConstraintKind (applyTys k args)
+    go_k k args = isConstraintKind (piResultTys k args)
 
 isClassPred, isEqPred, isNomEqPred, isIPPred :: PredType -> Bool
 isClassPred ty = case tyConAppTyCon_maybe ty of