Use a pattern synonym for arguments to FunTy (#18750)wip/funtycon-args

This makes it easier to add more arguments to FunTy in the future.

Not done in Unify.hs, because of perf problems #20165.

8 files changed, 72 insertions, 39 deletions

diff --git a/compiler/GHC/Core/Coercion.hs b/compiler/GHC/Core/Coercion.hs
index aa0cf29754..b6f501958e 100644
--- a/compiler/GHC/Core/Coercion.hs
+++ b/compiler/GHC/Core/Coercion.hs
@@ -1,6 +1,7 @@
 {-# LANGUAGE FlexibleContexts    #-}
 {-# LANGUAGE RankNTypes          #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE MultiWayIf          #-}
 
 {-
 (c) The University of Glasgow 2006
@@ -393,14 +394,15 @@ isReflMCo _     = False
 Note [Function coercions]
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 Remember that
-  (->) :: forall {r1} {r2}. TYPE r1 -> TYPE r2 -> TYPE LiftedRep
+  FUN :: forall (m :: Multiplicity) ->
+         forall {r1} {r2}. TYPE r1 -> TYPE r2 -> TYPE LiftedRep
 whose `RuntimeRep' arguments are intentionally marked inferred to
 avoid type application.
 
 Hence
-  FunCo r mult co1 co2 :: (s1->t1) ~r (s2->t2)
+  FunCo r co_mult co1 co2 :: (s1%m1->t1) ~r (s2%m2->t2)
 is short for
-  TyConAppCo (->) mult co_rep1 co_rep2 co1 co2
+  TyConAppCo FUN co_mult co_rep1 co_rep2 co1 co2
 where co_rep1, co_rep2 are the coercions on the representations.
 -}
 
@@ -430,7 +432,9 @@ decomposeFunCo _ (FunCo _ w co1 co2) = (w, co1, co2)
    -- Short-circuits the calls to mkNthCo
 
 decomposeFunCo r co = assertPpr all_ok (ppr co)
-                      (mkNthCo Nominal 0 co, mkNthCo r 3 co, mkNthCo r 4 co)
+                      (mkNthCo Nominal funTyConMulArgNo co,
+                       mkNthCo r       funTyConArgArgNo co,
+                       mkNthCo r       funTyConResArgNo co)
   where
     Pair s1t1 s2t2 = coercionKind co
     all_ok = isFunTy s1t1 && isFunTy s2t2
@@ -536,8 +540,8 @@ splitTyConAppCo_maybe co
        ; let args = zipWith mkReflCo (tyConRolesX r tc) tys
        ; return (tc, args) }
 splitTyConAppCo_maybe (TyConAppCo _ tc cos) = Just (tc, cos)
-splitTyConAppCo_maybe (FunCo _ w arg res)     = Just (funTyCon, cos)
-  where cos = [w, mkRuntimeRepCo arg, mkRuntimeRepCo res, arg, res]
+splitTyConAppCo_maybe (FunCo _ m arg res)     = Just (funTyCon, cos)
+  where cos = FunTyConArgs m (mkRuntimeRepCo arg) (mkRuntimeRepCo res) arg res
 splitTyConAppCo_maybe _                     = Nothing
 
 multToCo :: Mult -> Coercion
@@ -775,6 +779,7 @@ mkNomReflCo = Refl
 mkTyConAppCo :: HasDebugCallStack => Role -> TyCon -> [Coercion] -> Coercion
 mkTyConAppCo r tc cos
   | [w, _rep1, _rep2, co1, co2] <- cos   -- See Note [Function coercions]
+    -- We should use FunTyConArgs here, but this causes perf problems (#20165)
   , isFunTyCon tc
   = -- (a :: TYPE ra) -> (b :: TYPE rb)  ~  (c :: TYPE rc) -> (d :: TYPE rd)
     -- rep1 :: ra  ~  rc        rep2 :: rb  ~  rd
@@ -1175,21 +1180,21 @@ mkNthCoFunCo :: Int         -- ^ "n"
              -> Coercion    -- ^ result coercion
              -> Coercion    -- ^ nth coercion from a FunCo
 -- See Note [Function coercions]
--- If FunCo _ mult arg_co res_co ::   (s1:TYPE sk1 :mult-> s2:TYPE sk2)
---                                  ~ (t1:TYPE tk1 :mult-> t2:TYPE tk2)
+-- If FunCo _ mult_co arg_co res_co ::   (s1:TYPE sk1 %mult1-> s2:TYPE sk2)
+--                                     ~ (t1:TYPE tk1 %mult2-> t2:TYPE tk2)
 -- Then we want to behave as if co was
---    TyConAppCo mult argk_co resk_co arg_co res_co
+--    TyConAppCo mult_co argk_co resk_co arg_co res_co
 -- where
 --    argk_co :: sk1 ~ tk1  =  mkNthCo 0 (mkKindCo arg_co)
 --    resk_co :: sk2 ~ tk2  =  mkNthCo 0 (mkKindCo res_co)
 --                             i.e. mkRuntimeRepCo
-mkNthCoFunCo n w co1 co2 = case n of
-  0 -> w
-  1 -> mkRuntimeRepCo co1
-  2 -> mkRuntimeRepCo co2
-  3 -> co1
-  4 -> co2
-  _ -> pprPanic "mkNthCo(FunCo)" (ppr n $$ ppr w $$ ppr co1 $$ ppr co2)
+mkNthCoFunCo n com co1 co2 =
+  if | n == funTyConMulArgNo   -> com
+     | n == funTyConArgRRArgNo -> mkRuntimeRepCo co1
+     | n == funTyConResRRArgNo -> mkRuntimeRepCo co2
+     | n == funTyConArgArgNo   -> co1
+     | n == funTyConResArgNo   -> co2
+     | otherwise -> pprPanic "mkNthCo(FunCo)" (ppr n $$ ppr com $$ ppr co1 $$ ppr co2)
 
 -- | If you're about to call @mkNthCo r n co@, then @r@ should be
 -- whatever @nthCoRole n co@ returns.
@@ -1557,7 +1562,7 @@ instCoercion (Pair lty rty) g w
   | isFunTy lty && isFunTy rty
     -- g :: (t1 -> t2) ~ (t3 -> t4)
     -- returns t2 ~ t4
-  = Just $ mkNthCo Nominal 4 g -- extract result type, which is the 5th argument to (->)
+  = Just $ mkNthCo Nominal funTyConResArgNo g -- extract result type
   | otherwise -- one forall, one funty...
   = Nothing
 
diff --git a/compiler/GHC/Core/Lint.hs b/compiler/GHC/Core/Lint.hs
index f20dbcc62b..02e4791196 100644
--- a/compiler/GHC/Core/Lint.hs
+++ b/compiler/GHC/Core/Lint.hs
@@ -1632,7 +1632,7 @@ lintType ty@(TyConApp tc tys)
        ; lintTySynFamApp report_unsat ty tc tys }
 
   | isFunTyCon tc
-  , tys `lengthIs` 5
+  , FunTyConArgs _ _ _ _ _ <- tys
     -- We should never see a saturated application of funTyCon; such
     -- applications should be represented with the FunTy constructor.
     -- See Note [Linting function types] and
@@ -2000,7 +2000,7 @@ lintCoercion (GRefl r ty (MCo co))
 
 lintCoercion co@(TyConAppCo r tc cos)
   | tc `hasKey` funTyConKey
-  , [_w, _rep1,_rep2,_co1,_co2] <- cos
+  , FunTyConArgs _w _rep1 _rep2 _co1 _co2 <- cos
   = failWithL (text "Saturated TyConAppCo (->):" <+> ppr co)
     -- All saturated TyConAppCos should be FunCos
 
diff --git a/compiler/GHC/Core/TyCon.hs b/compiler/GHC/Core/TyCon.hs
index 512ac4737c..e24b58fdd3 100644
--- a/compiler/GHC/Core/TyCon.hs
+++ b/compiler/GHC/Core/TyCon.hs
@@ -11,6 +11,8 @@
 The @TyCon@ datatype
 -}
 
+{-# LANGUAGE PatternSynonyms #-}
+
 module GHC.Core.TyCon(
         -- * Main TyCon data types
         TyCon,
@@ -30,6 +32,13 @@ module GHC.Core.TyCon(
         -- ** Field labels
         tyConFieldLabels, lookupTyConFieldLabel,
 
+        pattern FunTyConArgs,
+        funTyConMulArgNo,
+        funTyConArgRRArgNo,
+        funTyConResRRArgNo,
+        funTyConArgArgNo,
+        funTyConResArgNo,
+
         -- ** Constructing TyCons
         mkAlgTyCon,
         mkClassTyCon,
@@ -1705,6 +1714,24 @@ module mutual-recursion.  And they aren't called from many places.
 So we compromise, and move their Kind calculation to the call site.
 -}
 
+-- | This pattern synonym signifies arguments to the FUN type,
+-- allowing easier grepping (#18750).
+-- If it cannot be used in some place, you should leave a comment
+-- mentioning FunTyConArgs.
+-- The same construct is used for FunCos, see Note [Function coercions]
+-- in GHC.Core.Coercion.
+pattern FunTyConArgs :: a -> a -> a -> a -> a -> [a]
+pattern FunTyConArgs mul r1 r2 arg res = [mul, r1, r2, arg, res]
+
+-- Integers denoting positions of arguments to FunTy.
+funTyConMulArgNo, funTyConArgRRArgNo, funTyConResRRArgNo, funTyConArgArgNo,
+  funTyConResArgNo :: Int
+funTyConMulArgNo   = 0
+funTyConArgRRArgNo = 1
+funTyConResRRArgNo = 2
+funTyConArgArgNo   = 3
+funTyConResArgNo   = 4
+
 -- | Given the name of the function type constructor and it's kind, create the
 -- corresponding 'TyCon'. It is recommended to use 'GHC.Core.TyCo.Rep.funTyCon' if you want
 -- this functionality
@@ -2516,7 +2543,7 @@ tyConRoles :: TyCon -> [Role]
 -- See also Note [TyCon Role signatures]
 tyConRoles tc
   = case tc of
-    { FunTyCon {}                         -> [Nominal, Nominal, Nominal, Representational, Representational]
+    { FunTyCon {}                         -> FunTyConArgs Nominal Nominal Nominal Representational Representational
     ; AlgTyCon { tcRoles = roles }        -> roles
     ; SynonymTyCon { tcRoles = roles }    -> roles
     ; FamilyTyCon {}                      -> const_role Nominal
diff --git a/compiler/GHC/Core/Type.hs b/compiler/GHC/Core/Type.hs
index 6b88262ff5..5e06ae3a1c 100644
--- a/compiler/GHC/Core/Type.hs
+++ b/compiler/GHC/Core/Type.hs
@@ -1036,7 +1036,7 @@ repSplitAppTy_maybe :: HasDebugCallStack => Type -> Maybe (Type,Type)
 -- ^ Does the AppTy split as in 'splitAppTy_maybe', but assumes that
 -- any Core view stuff is already done
 repSplitAppTy_maybe (FunTy _ w ty1 ty2)
-  = Just (TyConApp funTyCon [w, rep1, rep2, ty1], ty2)
+  = Just (TyConApp funTyCon [w, rep1, rep2, ty1], ty2)  -- can't use FunTyConArgs - the last argument is removed
   where
     rep1 = getRuntimeRep ty1
     rep2 = getRuntimeRep ty2
@@ -1061,7 +1061,7 @@ tcRepSplitAppTy_maybe (FunTy { ft_af = af, ft_mult = w, ft_arg = ty1, ft_res = t
   | InvisArg <- af
   = Nothing  -- See Note [Decomposing fat arrow c=>t]
   | otherwise
-  = Just (TyConApp funTyCon [w, rep1, rep2, ty1], ty2)
+  = Just (TyConApp funTyCon [w, rep1, rep2, ty1], ty2) -- can't use FunTyConArgs - the last argument is removed
   where
     rep1 = getRuntimeRep ty1
     rep2 = getRuntimeRep ty2
@@ -1099,7 +1099,7 @@ splitAppTys ty = split ty ty []
         (TyConApp tc tc_args1, tc_args2 ++ args)
     split _   (FunTy _ w ty1 ty2) args
       = assert (null args )
-        (TyConApp funTyCon [], [w, rep1, rep2, ty1, ty2])
+        (TyConApp funTyCon [], FunTyConArgs w rep1 rep2 ty1 ty2)
       where
         rep1 = getRuntimeRep ty1
         rep2 = getRuntimeRep ty2
@@ -1119,7 +1119,7 @@ repSplitAppTys ty = split ty []
         (TyConApp tc tc_args1, tc_args2 ++ args)
     split (FunTy _ w ty1 ty2) args
       = assert (null args )
-        (TyConApp funTyCon [], [w, rep1, rep2, ty1, ty2])
+        (TyConApp funTyCon [], FunTyConArgs w rep1 rep2 ty1 ty2)
       where
         rep1 = getRuntimeRep ty1
         rep2 = getRuntimeRep ty2
@@ -1433,7 +1433,7 @@ tyConAppArgs_maybe ty = case coreFullView ty of
   FunTy _ w arg res
     | Just rep1 <- getRuntimeRep_maybe arg
     , Just rep2 <- getRuntimeRep_maybe res
-    -> Just [w, rep1, rep2, arg, res]
+    -> Just (FunTyConArgs w rep1 rep2 arg res)
   _ -> Nothing
 
 tyConAppArgs :: Type -> [Type]
@@ -1494,7 +1494,7 @@ repSplitTyConApp_maybe (TyConApp tc tys) = Just (tc, tys)
 repSplitTyConApp_maybe (FunTy _ w arg res)
   | Just arg_rep <- getRuntimeRep_maybe arg
   , Just res_rep <- getRuntimeRep_maybe res
-  = Just (funTyCon, [w, arg_rep, res_rep, arg, res])
+  = Just (funTyCon, FunTyConArgs w arg_rep res_rep arg res)
 repSplitTyConApp_maybe _ = Nothing
 
 -------------------
@@ -1608,7 +1608,7 @@ mkTyConApp tycon tys
   = mkTyConTy tycon
 
   | isFunTyCon tycon
-  , [w, _rep1,_rep2,ty1,ty2] <- tys
+  , FunTyConArgs w _rep1 _rep2 ty1 ty2 <- tys
   -- The FunTyCon (->) is always a visible one
   = FunTy { ft_af = VisArg, ft_mult = w, ft_arg = ty1, ft_res = ty2 }
 
diff --git a/compiler/GHC/Core/Unify.hs b/compiler/GHC/Core/Unify.hs
index bd54ecee39..03394eb634 100644
--- a/compiler/GHC/Core/Unify.hs
+++ b/compiler/GHC/Core/Unify.hs
@@ -1611,7 +1611,7 @@ ty_co_match menv subst (FunTy _ w ty1 ty2) co _lkco _rkco
     -- runtime rep, a multiplicity and two types), we shouldn't need to
     -- explicitly unify the runtime reps here; unifying the types themselves
     -- should be sufficient.  See Note [Representation of function types].
-  | Just (tc, [co_mult, _,_,co1,co2]) <- splitTyConAppCo_maybe co
+  | Just (tc, FunTyConArgs co_mult _ _ co1 co2) <- splitTyConAppCo_maybe co
   , tc == funTyCon
   = let Pair lkcos rkcos = traverse (fmap mkNomReflCo . coercionKind) [co_mult,co1,co2]
     in ty_co_match_args menv subst [w, ty1, ty2] [co_mult, co1, co2] lkcos rkcos
@@ -1718,11 +1718,12 @@ pushRefl co =
   case (isReflCo_maybe co) of
     Just (AppTy ty1 ty2, Nominal)
       -> Just (AppCo (mkReflCo Nominal ty1) (mkNomReflCo ty2))
-    Just (FunTy _ w ty1 ty2, r)
+    Just (FunTy _ m ty1 ty2, r)
       | Just rep1 <- getRuntimeRep_maybe ty1
       , Just rep2 <- getRuntimeRep_maybe ty2
-      ->  Just (TyConAppCo r funTyCon [ multToCo w, mkReflCo r rep1, mkReflCo r rep2
-                                       , mkReflCo r ty1,  mkReflCo r ty2 ])
+      ->  Just (TyConAppCo r funTyCon (FunTyConArgs (multToCo m)
+                                                    (mkReflCo r rep1) (mkReflCo r rep2)
+                                                    (mkReflCo r ty1)  (mkReflCo r ty2)))
     Just (TyConApp tc tys, r)
       -> Just (TyConAppCo r tc (zipWith mkReflCo (tyConRolesX r tc) tys))
     Just (ForAllTy (Bndr tv _) ty, r)
diff --git a/compiler/GHC/CoreToIface.hs b/compiler/GHC/CoreToIface.hs
index 2356f6c7f5..53938a801e 100644
--- a/compiler/GHC/CoreToIface.hs
+++ b/compiler/GHC/CoreToIface.hs
@@ -302,8 +302,8 @@ toIfaceCoercionX fr co
                                           (toIfaceTypeX fr t2)
     go (TyConAppCo r tc cos)
       | tc `hasKey` funTyConKey
-      , [_,_,_,_, _] <- cos         = panic "toIfaceCoercion"
-      | otherwise                =
+      , FunTyConArgs _ _ _ _ _  <- cos = panic "toIfaceCoercion"
+      | otherwise                      =
         IfaceTyConAppCo r (toIfaceTyCon tc) (map go cos)
     go (FunCo r w co1 co2)   = IfaceFunCo r (go w) (go co1) (go co2)
 
diff --git a/compiler/GHC/Tc/Solver/Canonical.hs b/compiler/GHC/Tc/Solver/Canonical.hs
index 162ef60cbc..40769aa4b6 100644
--- a/compiler/GHC/Tc/Solver/Canonical.hs
+++ b/compiler/GHC/Tc/Solver/Canonical.hs
@@ -1023,8 +1023,8 @@ can_eq_nc' _rewritten _rdr_env _envs ev eq_rel
   , Just ty2a_rep <- getRuntimeRep_maybe ty2a
   , Just ty2b_rep <- getRuntimeRep_maybe ty2b
   = canDecomposableTyConAppOK ev eq_rel funTyCon
-                              [am1, ty1a_rep, ty1b_rep, ty1a, ty1b]
-                              [am2, ty2a_rep, ty2b_rep, ty2a, ty2b]
+                              (FunTyConArgs am1 ty1a_rep ty1b_rep ty1a ty1b)
+                              (FunTyConArgs am2 ty2a_rep ty2b_rep ty2a ty2b)
 
 -- Decompose type constructor applications
 -- NB: we have expanded type synonyms already
diff --git a/compiler/GHC/Tc/Utils/TcType.hs b/compiler/GHC/Tc/Utils/TcType.hs
index 367922e3e5..b5dd9f89cf 100644
--- a/compiler/GHC/Tc/Utils/TcType.hs
+++ b/compiler/GHC/Tc/Utils/TcType.hs
@@ -1658,7 +1658,7 @@ tc_eq_type keep_syns vis_only orig_ty1 orig_ty2
 
     orig_env = mkRnEnv2 $ mkInScopeSet $ tyCoVarsOfTypes [orig_ty1, orig_ty2]
 
-    -- @eqFunTy w arg res ty@ is True when @ty@ equals @FunTy w arg res@. This is
+    -- @eqFunTy m arg res ty@ is True when @ty@ equals @FunTy m arg res@. This is
     -- sometimes hard to know directly because @ty@ might have some casts
     -- obscuring the FunTy. And 'splitAppTy' is difficult because we can't
     -- always extract a RuntimeRep (see Note [xyz]) if the kind of the arg or
@@ -1666,15 +1666,15 @@ tc_eq_type keep_syns vis_only orig_ty1 orig_ty2
     -- corner case.
     eqFunTy :: RnEnv2 -> Mult -> Type -> Type -> Type -> Bool
                -- Last arg is /not/ FunTy
-    eqFunTy env w arg res ty@(AppTy{}) = get_args ty []
+    eqFunTy env m arg res ty@(AppTy{}) = get_args ty []
       where
         get_args :: Type -> [Type] -> Bool
         get_args (AppTy f x)       args = get_args f (x:args)
         get_args (CastTy t _)      args = get_args t args
         get_args (TyConApp tc tys) args
           | tc == funTyCon
-          , [w', _, _, arg', res'] <- tys ++ args
-          = go env w w' && go env arg arg' && go env res res'
+          , FunTyConArgs m' _ _ arg' res' <- tys ++ args
+          = go env m m' && go env arg arg' && go env res res'
         get_args _ _    = False
     eqFunTy _ _ _ _ _   = False
 {-# INLINE tc_eq_type #-} -- See Note [Specialising tc_eq_type].