Cache coercion roles in NthCowip/tdammers/D4394-squash

Most callers of mkNthCo know the role of the coercion they
are trying to make. So instead of calculating this role, we
pass it in, forcing the caller to calculate it when needed.

This introduces a performance regression in perf/compiler/T9872d.

1 files changed, 12 insertions, 12 deletions

diff --git a/compiler/coreSyn/CoreOpt.hs b/compiler/coreSyn/CoreOpt.hs
index 04e604eb06..68b826cd7d 100644
--- a/compiler/coreSyn/CoreOpt.hs
+++ b/compiler/coreSyn/CoreOpt.hs
@@ -928,13 +928,13 @@ Here we implement the "push rules" from FC papers:
   by pushing the coercion into the arguments
 -}
 
-pushCoArgs :: Coercion -> [CoreArg] -> Maybe ([CoreArg], Coercion)
+pushCoArgs :: CoercionR -> [CoreArg] -> Maybe ([CoreArg], Coercion)
 pushCoArgs co []         = return ([], co)
 pushCoArgs co (arg:args) = do { (arg',  co1) <- pushCoArg  co  arg
                               ; (args', co2) <- pushCoArgs co1 args
                               ; return (arg':args', co2) }
 
-pushCoArg :: Coercion -> CoreArg -> Maybe (CoreArg, Coercion)
+pushCoArg :: CoercionR -> CoreArg -> Maybe (CoreArg, Coercion)
 -- We have (fun |> co) arg, and we want to transform it to
 --         (fun arg) |> co
 -- This may fail, e.g. if (fun :: N) where N is a newtype
@@ -967,7 +967,7 @@ pushCoTyArg co ty
        -- tyL = forall (a1 :: k1). ty1
        -- tyR = forall (a2 :: k2). ty2
 
-    co1 = mkNthCo 0 co
+    co1 = mkNthCo Nominal 0 co
        -- co1 :: k1 ~N k2
        -- Note that NthCo can extract a Nominal equality between the
        -- kinds of the types related by a coercion between forall-types.
@@ -977,7 +977,7 @@ pushCoTyArg co ty
         -- co2 :: ty1[ (ty|>co1)/a1 ] ~ ty2[ ty/a2 ]
         -- Arg of mkInstCo is always nominal, hence mkNomReflCo
 
-pushCoValArg :: Coercion -> Maybe (Coercion, Coercion)
+pushCoValArg :: CoercionR -> Maybe (Coercion, Coercion)
 -- We have (fun |> co) arg
 -- Push the coercion through to return
 --         (fun (arg |> co_arg)) |> co_res
@@ -987,7 +987,7 @@ pushCoValArg co
   = Just (mkRepReflCo arg, mkRepReflCo res)
 
   | isFunTy tyL
-  , (co1, co2) <- decomposeFunCo co
+  , (co1, co2) <- decomposeFunCo Representational co
               -- If   co  :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
               -- then co1 :: tyL1 ~ tyR1
               --      co2 :: tyL2 ~ tyR2
@@ -1001,7 +1001,7 @@ pushCoValArg co
     Pair tyL tyR = coercionKind co
 
 pushCoercionIntoLambda
-    :: InScopeSet -> Var -> CoreExpr -> Coercion -> Maybe (Var, CoreExpr)
+    :: InScopeSet -> Var -> CoreExpr -> CoercionR -> Maybe (Var, CoreExpr)
 -- This implements the Push rule from the paper on coercions
 --    (\x. e) |> co
 -- ===>
@@ -1011,7 +1011,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 (co1, co2) = decomposeFunCo co
+    = let (co1, co2) = decomposeFunCo Representational co
           -- Should we optimize the coercions here?
           -- Otherwise they might not match too well
           x' = x `setIdType` t1
@@ -1057,7 +1057,7 @@ pushCoDataCon dc dc_args co
         (ex_args, val_args) = splitAtList dc_ex_tyvars non_univ_args
 
         -- Make the "Psi" from the paper
-        omegas = decomposeCo tc_arity co
+        omegas = decomposeCo tc_arity co (tyConRolesRepresentational to_tc)
         (psi_subst, to_ex_arg_tys)
           = liftCoSubstWithEx Representational
                               dc_univ_tyvars
@@ -1104,7 +1104,7 @@ collectBindersPushingCo e
     go bs e           = (reverse bs, e)
 
     -- We are in a cast; peel off casts until we hit a lambda.
-    go_c :: [Var] -> CoreExpr -> Coercion -> ([Var], CoreExpr)
+    go_c :: [Var] -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
     -- (go_c bs e c) is same as (go bs e (e |> c))
     go_c bs (Cast e co1) co2 = go_c bs e (co1 `mkTransCo` co2)
     go_c bs (Lam b e)    co  = go_lam bs b e co
@@ -1112,20 +1112,20 @@ collectBindersPushingCo e
 
     -- We are in a lambda under a cast; peel off lambdas and build a
     -- new coercion for the body.
-    go_lam :: [Var] -> Var -> CoreExpr -> Coercion -> ([Var], CoreExpr)
+    go_lam :: [Var] -> Var -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
     -- (go_lam bs b e c) is same as (go_c bs (\b.e) c)
     go_lam bs b e co
       | isTyVar b
       , let Pair tyL tyR = coercionKind co
       , ASSERT( isForAllTy tyL )
         isForAllTy tyR
-      , isReflCo (mkNthCo 0 co)  -- See Note [collectBindersPushingCo]
+      , isReflCo (mkNthCo Nominal 0 co)  -- See Note [collectBindersPushingCo]
       = go_c (b:bs) e (mkInstCo co (mkNomReflCo (mkTyVarTy b)))
 
       | isId b
       , let Pair tyL tyR = coercionKind co
       , ASSERT( isFunTy tyL) isFunTy tyR
-      , (co_arg, co_res) <- decomposeFunCo co
+      , (co_arg, co_res) <- decomposeFunCo Representational co
       , isReflCo co_arg  -- See Note [collectBindersPushingCo]
       = go_c (b:bs) e co_res