Fix a nasty bug in CoreSubst.collectBindersPushingCowip/spj-early-inline2

The bug wsa in the use of (mkNthCo 0) to get the argument
part of a function coercion.  Not so! Now (->) takes four
arguments so that 0 should have been 2.

Enough with magic numbers. I defined decomposeFunCo, and used
it throughout.  Much nicer now; and correct.

The nete effect, incidentally, was that T9509 was failing to
specialise.  (And that was the initial reason for introducing
collectBindersPushingCo in the first place.)

3 files changed, 47 insertions, 10 deletions

diff --git a/compiler/coreSyn/CoreSubst.hs b/compiler/coreSyn/CoreSubst.hs
index 6afa3bab5a..5072e705f1 100644
--- a/compiler/coreSyn/CoreSubst.hs
+++ b/compiler/coreSyn/CoreSubst.hs
@@ -1678,7 +1678,7 @@ pushCoValArg co
   = Just (mkRepReflCo arg, mkRepReflCo res)
 
   | isFunTy tyL
-  , [_, _, co1, co2] <- decomposeCo 4 co
+  , (co1, co2) <- decomposeFunCo co
               -- If   co  :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
               -- then co1 :: tyL1 ~ tyR1
               --      co2 :: tyL2 ~ tyR2
@@ -1702,7 +1702,7 @@ pushCoercionIntoLambda in_scope x e co
     , Pair s1s2 t1t2 <- coercionKind co
     , Just (_s1,_s2) <- splitFunTy_maybe s1s2
     , Just (t1,_t2) <- splitFunTy_maybe t1t2
-    = let [_rep1, _rep2, co1, co2] = decomposeCo 4 co
+    = let (co1, co2) = decomposeFunCo co
           -- Should we optimize the coercions here?
           -- Otherwise they might not match too well
           x' = x `setIdType` t1
@@ -1808,8 +1808,9 @@ collectBindersPushingCo e
       | isId b
       , let Pair tyL tyR = coercionKind co
       , ASSERT( isFunTy tyL) isFunTy tyR
-      , isReflCo (mkNthCo 0 co)  -- See Note [collectBindersPushingCo]
-      = go_c (b:bs) e (mkNthCo 1 co)
+      , (co_arg, co_res) <- decomposeFunCo co
+      , isReflCo co_arg  -- See Note [collectBindersPushingCo]
+      = go_c (b:bs) e co_res
 
       | otherwise = (reverse bs, mkCast (Lam b e) co)
 
diff --git a/compiler/types/Coercion.hs b/compiler/types/Coercion.hs
index f2351fe0d3..f53968edc1 100644
--- a/compiler/types/Coercion.hs
+++ b/compiler/types/Coercion.hs
@@ -48,7 +48,7 @@ module Coercion (
         mapStepResult, unwrapNewTypeStepper,
         topNormaliseNewType_maybe, topNormaliseTypeX,
 
-        decomposeCo, getCoVar_maybe,
+        decomposeCo, decomposeFunCo, getCoVar_maybe,
         splitTyConAppCo_maybe,
         splitAppCo_maybe,
         splitFunCo_maybe,
@@ -296,8 +296,20 @@ ppr_co_ax_branch ppr_rhs
         Destructing coercions
 %*                                                                      *
 %************************************************************************
+
+Note [Function coercions]
+~~~~~~~~~~~~~~~~~~~~~~~~~
+Remember that
+  (->) :: forall r1 r2. TYPE r1 -> TYPE r2 -> TYPE LiftedRep
+
+Hence
+  FunCo r co1 co2 :: (s1->t1) ~r (s2->t2)
+is short for
+  TyConAppCo (->) co_rep1 co_rep2 co1 co2
+where co_rep1, co_rep2 are the coercions on the representations.
 -}
 
+
 -- | This breaks a 'Coercion' with type @T A B C ~ T D E F@ into
 -- a list of 'Coercion's of kinds @A ~ D@, @B ~ E@ and @E ~ F@. Hence:
 --
@@ -307,6 +319,16 @@ decomposeCo arity co
   = [mkNthCo n co | n <- [0..(arity-1)] ]
            -- Remember, Nth is zero-indexed
 
+decomposeFunCo :: Coercion -> (Coercion, Coercion)
+-- Expects co :: (s1 -> t1) ~ (s2 -> t2)
+-- Returns (co1 :: s1~s2, co2 :: t1~t2)
+-- See Note [Function coercions] for the "2" and "3"
+decomposeFunCo co = ASSERT2( all_ok, ppr co )
+                    (mkNthCo 2 co, mkNthCo 3 co)
+  where
+    Pair s1t1 s2t2 = coercionKind co
+    all_ok = isFunTy s1t1 && isFunTy s2t2
+
 -- | Attempts to obtain the type variable underlying a 'Coercion'
 getCoVar_maybe :: Coercion -> Maybe CoVar
 getCoVar_maybe (CoVarCo cv) = Just cv
@@ -595,7 +617,7 @@ mkNomReflCo = mkReflCo Nominal
 mkTyConAppCo :: HasDebugCallStack => Role -> TyCon -> [Coercion] -> Coercion
 mkTyConAppCo r tc cos
   | tc `hasKey` funTyConKey
-  , [_rep1, _rep2, co1, co2] <- cos
+  , [_rep1, _rep2, co1, co2] <- cos   -- See Note [Function coercions]
   = -- (a :: TYPE ra) -> (b :: TYPE rb)  ~  (c :: TYPE rc) -> (d :: TYPE rd)
     -- rep1 :: ra  ~  rc        rep2 :: rb  ~  rd
     -- co1  :: a   ~  c         co2  :: b   ~  d
@@ -923,14 +945,26 @@ mkNthCo n (Refl r ty)
 mkNthCo 0 (ForAllCo _ kind_co _) = kind_co
   -- If co :: (forall a1:k1. t1) ~ (forall a2:k2. t2)
   -- then (nth 0 co :: k1 ~ k2)
-mkNthCo n (TyConAppCo _ _ arg_cos) = arg_cos `getNth` n
+
 mkNthCo n co@(FunCo _ arg res)
+  -- See Note [Function coercions]
+  -- If FunCo _ arg_co res_co ::   (s1:TYPE sk1 -> s2:TYPE sk2)
+  --                             ~ (t1:TYPE tk1 -> t2:TYPE tk2)
+  -- Then we want to behave as if co was
+  --    TyConAppCo 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 -> mkRuntimeRepCo arg
       1 -> mkRuntimeRepCo res
       2 -> arg
       3 -> res
       _ -> pprPanic "mkNthCo(FunCo)" (ppr n $$ ppr co)
+
+mkNthCo n (TyConAppCo _ _ arg_cos) = arg_cos `getNth` n
+
 mkNthCo n co = NthCo n co
 
 mkLRCo :: LeftOrRight -> Coercion -> Coercion
@@ -978,8 +1012,10 @@ mkKindCo co
        -- generally, calling coercionKind during coercion creation is a bad idea,
        -- as it can lead to exponential behavior. But, we don't have nested mkKindCos,
        -- so it's OK here.
-  , typeKind ty1 `eqType` typeKind ty2
-  = Refl Nominal (typeKind ty1)
+  , let tk1 = typeKind ty1
+        tk2 = typeKind ty2
+  , tk1 `eqType` tk2
+  = Refl Nominal tk1
   | otherwise
   = KindCo co
 
diff --git a/compiler/types/Unify.hs b/compiler/types/Unify.hs
index 05d6c6d091..517358d482 100644
--- a/compiler/types/Unify.hs
+++ b/compiler/types/Unify.hs
@@ -843,7 +843,7 @@ unify_ty (CoercionTy co1) (CoercionTy co2) kco
              -> do { b <- tvBindFlagL cv
                    ; if b == BindMe
                        then do { checkRnEnvRCo co2
-                               ; let [_, _, co_l, co_r] = decomposeCo 4 kco
+                               ; let (co_l, co_r) = decomposeFunCo kco
                                   -- cv :: t1 ~ t2
                                   -- co2 :: s1 ~ s2
                                   -- co_l :: t1 ~ s1