Optimize NthCo (FunCo ...) in coercion opt

We were missing this case previously.

Close #18528.

Metric Decrease:
    T18223
    T5321Fun

3 files changed, 39 insertions, 28 deletions

diff --git a/compiler/GHC/Builtin/Types/Prim.hs b/compiler/GHC/Builtin/Types/Prim.hs
index e631e04baa..e930b58b36 100644
--- a/compiler/GHC/Builtin/Types/Prim.hs
+++ b/compiler/GHC/Builtin/Types/Prim.hs
@@ -407,8 +407,9 @@ funTyConName = mkPrimTyConName (fsLit "FUN") funTyConKey funTyCon
 -- | The @FUN@ type constructor.
 --
 -- @
--- FUN :: forall {m :: Multiplicity} {rep1 :: RuntimeRep} {rep2 :: RuntimeRep}.
---         TYPE rep1 -> TYPE rep2 -> *
+-- FUN :: forall (m :: Multiplicity) ->
+--        forall {rep1 :: RuntimeRep} {rep2 :: RuntimeRep}.
+--        TYPE rep1 -> TYPE rep2 -> *
 -- @
 --
 -- The runtime representations quantification is left inferred. This
diff --git a/compiler/GHC/Core/Coercion.hs b/compiler/GHC/Core/Coercion.hs
index 6ba6c0f144..647084baff 100644
--- a/compiler/GHC/Core/Coercion.hs
+++ b/compiler/GHC/Core/Coercion.hs
@@ -32,7 +32,7 @@ module GHC.Core.Coercion (
         mkAxInstLHS, mkUnbranchedAxInstLHS,
         mkPiCo, mkPiCos, mkCoCast,
         mkSymCo, mkTransCo,
-        mkNthCo, nthCoRole, mkLRCo,
+        mkNthCo, mkNthCoFunCo, nthCoRole, mkLRCo,
         mkInstCo, mkAppCo, mkAppCos, mkTyConAppCo, mkFunCo,
         mkForAllCo, mkForAllCos, mkHomoForAllCos,
         mkPhantomCo,
@@ -1052,23 +1052,8 @@ mkNthCo r n co
       -- If co :: (forall a1:t1 ~ t2. t1) ~ (forall a2:t3 ~ t4. t2)
       -- then (nth 0 co :: (t1 ~ t2) ~N (t3 ~ t4))
 
-    go r n co@(FunCo r0 w arg res)
-      -- 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)
-      -- Then we want to behave as if co was
-      --    TyConAppCo mult 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
-      = case n of
-          0 -> ASSERT( r == Nominal ) w
-          1 -> ASSERT( r == Nominal ) mkRuntimeRepCo arg
-          2 -> ASSERT( r == Nominal ) mkRuntimeRepCo res
-          3 -> ASSERT( r == r0 )      arg
-          4 -> ASSERT( r == r0 )      res
-          _ -> pprPanic "mkNthCo(FunCo)" (ppr n $$ ppr co)
+    go _ n (FunCo _ w arg res)
+      = mkNthCoFunCo n w arg res
 
     go r n (TyConAppCo r0 tc arg_cos) = ASSERT2( r == nthRole r0 tc n
                                                     , (vcat [ ppr tc
@@ -1120,7 +1105,28 @@ mkNthCo r n co
       | otherwise
       = True
 
-
+-- | Extract the nth field of a FunCo
+mkNthCoFunCo :: Int         -- ^ "n"
+             -> CoercionN   -- ^ multiplicity coercion
+             -> Coercion    -- ^ argument coercion
+             -> 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)
+-- Then we want to behave as if co was
+--    TyConAppCo mult 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)
 
 -- | If you're about to call @mkNthCo r n co@, then @r@ should be
 -- whatever @nthCoRole n co@ returns.
diff --git a/compiler/GHC/Core/Coercion/Opt.hs b/compiler/GHC/Core/Coercion/Opt.hs
index b9656a45bb..fb0a6b0cc0 100644
--- a/compiler/GHC/Core/Coercion/Opt.hs
+++ b/compiler/GHC/Core/Coercion/Opt.hs
@@ -332,6 +332,7 @@ opt_co4 env _sym rep r (NthCo _r n co)
   , Just (_tc, args) <- ASSERT( r == _r )
                         splitTyConApp_maybe ty
   = liftCoSubst (chooseRole rep r) env (args `getNth` n)
+
   | Just (ty, _) <- isReflCo_maybe co
   , n == 0
   , Just (tv, _) <- splitForAllTy_maybe ty
@@ -342,6 +343,11 @@ opt_co4 env sym rep r (NthCo r1 n (TyConAppCo _ _ cos))
   = ASSERT( r == r1 )
     opt_co4_wrap env sym rep r (cos `getNth` n)
 
+-- see the definition of GHC.Builtin.Types.Prim.funTyCon
+opt_co4 env sym rep r (NthCo r1 n (FunCo _r2 w co1 co2))
+  = ASSERT( r == r1 )
+    opt_co4_wrap env sym rep r (mkNthCoFunCo n w co1 co2)
+
 opt_co4 env sym rep r (NthCo _r n (ForAllCo _ eta _))
       -- works for both tyvar and covar
   = ASSERT( r == _r )
@@ -349,18 +355,16 @@ opt_co4 env sym rep r (NthCo _r n (ForAllCo _ eta _))
     opt_co4_wrap env sym rep Nominal eta
 
 opt_co4 env sym rep r (NthCo _r n co)
-  | TyConAppCo _ _ cos <- co'
-  , let nth_co = cos `getNth` n
+  | Just nth_co <- case co' of
+      TyConAppCo _ _ cos -> Just (cos `getNth` n)
+      FunCo _ w co1 co2  -> Just (mkNthCoFunCo n w co1 co2)
+      ForAllCo _ eta _   -> Just eta
+      _                  -> Nothing
   = if rep && (r == Nominal)
       -- keep propagating the SubCo
     then opt_co4_wrap (zapLiftingContext env) False True Nominal nth_co
     else nth_co
 
-  | ForAllCo _ eta _ <- co'
-  = if rep
-    then opt_co4_wrap (zapLiftingContext env) False True Nominal eta
-    else eta
-
   | otherwise
   = wrapRole rep r $ NthCo r n co'
   where