Caching coercion roles in NthCo and coercionKindsRole refactoring

While addressing nonlinear behavior related to coercion roles,
particularly `NthCo`, we noticed that coercion roles are recalculated
often even though they should be readily at hand already in most cases.
This patch adds a `Role` to the `NthCo` constructor so that we can cache
them rather than having to recalculate them on the fly.
https://ghc.haskell.org/trac/ghc/ticket/11735#comment:23 explains the
approach.

Performance improvement over GHC HEAD, when compiling Grammar.hs (see below):

GHC 8.2.1:
```
ghc Grammar.hs  176.27s user 0.23s system 99% cpu 2:56.81 total
```

before patch (but with other optimizations applied):
```
ghc Grammar.hs -fforce-recomp  175.77s user 0.19s system 100% cpu 2:55.78 total
```

after:
```
../../ghc/inplace/bin/ghc-stage2 Grammar.hs  10.32s user 0.17s system 98% cpu 10.678 total
```

Introduces the following regressions:

- perf/compiler/parsing001 (possibly false positive)
- perf/compiler/T9872
- perf/compiler/haddock.base

Reviewers: goldfire, bgamari, simonpj

Reviewed By: simonpj

Subscribers: rwbarton, thomie, carter

GHC Trac Issues: #11735

Differential Revision: https://phabricator.haskell.org/D4394

4 files changed, 19 insertions, 17 deletions

diff --git a/compiler/coreSyn/CoreFVs.hs b/compiler/coreSyn/CoreFVs.hs
index 7fcff90c71..7e7727164c 100644
--- a/compiler/coreSyn/CoreFVs.hs
+++ b/compiler/coreSyn/CoreFVs.hs
@@ -379,7 +379,7 @@ orphNamesOfCo (AxiomInstCo con _ cos) = orphNamesOfCoCon con `unionNameSet` orph
 orphNamesOfCo (UnivCo p _ t1 t2)    = orphNamesOfProv p `unionNameSet` orphNamesOfType t1 `unionNameSet` orphNamesOfType t2
 orphNamesOfCo (SymCo co)            = orphNamesOfCo co
 orphNamesOfCo (TransCo co1 co2)     = orphNamesOfCo co1 `unionNameSet` orphNamesOfCo co2
-orphNamesOfCo (NthCo _ co)          = orphNamesOfCo co
+orphNamesOfCo (NthCo _ _ co)        = orphNamesOfCo co
 orphNamesOfCo (LRCo  _ co)          = orphNamesOfCo co
 orphNamesOfCo (InstCo co arg)       = orphNamesOfCo co `unionNameSet` orphNamesOfCo arg
 orphNamesOfCo (CoherenceCo co1 co2) = orphNamesOfCo co1 `unionNameSet` orphNamesOfCo co2
diff --git a/compiler/coreSyn/CoreLint.hs b/compiler/coreSyn/CoreLint.hs
index 812e1ed281..e8319aed33 100644
--- a/compiler/coreSyn/CoreLint.hs
+++ b/compiler/coreSyn/CoreLint.hs
@@ -1759,12 +1759,13 @@ lintCoercion co@(TransCo co1 co2)
        ; lintRole co r1 r2
        ; return (k1a, k2b, ty1a, ty2b, r1) }
 
-lintCoercion the_co@(NthCo n co)
+lintCoercion the_co@(NthCo r0 n co)
   = do { (_, _, s, t, r) <- lintCoercion co
        ; case (splitForAllTy_maybe s, splitForAllTy_maybe t) of
          { (Just (tv_s, _ty_s), Just (tv_t, _ty_t))
              |  n == 0
-             -> return (ks, kt, ts, tt, Nominal)
+             -> do { lintRole the_co Nominal r0
+                   ; return (ks, kt, ts, tt, r0) }
              where
                ts = tyVarKind tv_s
                tt = tyVarKind tv_t
@@ -1778,7 +1779,8 @@ lintCoercion the_co@(NthCo n co)
                  -- see Note [NthCo and newtypes] in TyCoRep
              , tys_s `equalLength` tys_t
              , tys_s `lengthExceeds` n
-             -> return (ks, kt, ts, tt, tr)
+             -> do { lintRole the_co tr r0
+                   ; return (ks, kt, ts, tt, r0) }
              where
                ts = getNth tys_s n
                tt = getNth tys_t n
diff --git a/compiler/coreSyn/CoreOpt.hs b/compiler/coreSyn/CoreOpt.hs
index 42cc706bff..372c2ea5aa 100644
--- a/compiler/coreSyn/CoreOpt.hs
+++ b/compiler/coreSyn/CoreOpt.hs
@@ -930,7 +930,7 @@ Here we implement the "push rules" from FC papers:
   by pushing the coercion into the arguments
 -}
 
-pushCoArgs :: Coercion -> [CoreArg] -> Maybe ([CoreArg], Maybe Coercion)
+pushCoArgs :: CoercionR -> [CoreArg] -> Maybe ([CoreArg], Maybe Coercion)
 pushCoArgs co []         = return ([], Just co)
 pushCoArgs co (arg:args) = do { (arg',  m_co1) <- pushCoArg  co  arg
                               ; case m_co1 of
@@ -938,7 +938,7 @@ pushCoArgs co (arg:args) = do { (arg',  m_co1) <- pushCoArg  co  arg
                                                  ; return (arg':args', m_co2) }
                                   Nothing  -> return (arg':args, Nothing) }
 
-pushCoArg :: Coercion -> CoreArg -> Maybe (CoreArg, Maybe Coercion)
+pushCoArg :: CoercionR -> CoreArg -> Maybe (CoreArg, Maybe 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
@@ -973,7 +973,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.
@@ -983,7 +983,7 @@ pushCoTyArg co ty
         -- co2 :: ty1[ (ty|>co1)/a1 ] ~ ty2[ ty/a2 ]
         -- Arg of mkInstCo is always nominal, hence mkNomReflCo
 
-pushCoValArg :: Coercion -> Maybe (Coercion, Maybe Coercion)
+pushCoValArg :: CoercionR -> Maybe (Coercion, Maybe Coercion)
 -- We have (fun |> co) arg
 -- Push the coercion through to return
 --         (fun (arg |> co_arg)) |> co_res
@@ -995,7 +995,7 @@ pushCoValArg co
   = Just (mkRepReflCo arg, Nothing)
 
   | isFunTy tyL
-  , (co1, co2) <- decomposeFunCo co
+  , (co1, co2) <- decomposeFunCo Representational co
               -- If   co  :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
               -- then co1 :: tyL1 ~ tyR1
               --      co2 :: tyL2 ~ tyR2
@@ -1009,7 +1009,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
 -- ===>
@@ -1019,7 +1019,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
@@ -1065,7 +1065,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
@@ -1112,7 +1112,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
@@ -1120,20 +1120,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
 
diff --git a/compiler/coreSyn/CoreUtils.hs b/compiler/coreSyn/CoreUtils.hs
index 3c5cb3455e..4137c1c850 100644
--- a/compiler/coreSyn/CoreUtils.hs
+++ b/compiler/coreSyn/CoreUtils.hs
@@ -257,7 +257,7 @@ applyTypeToArgs e op_ty args
 
 -- | Wrap the given expression in the coercion safely, dropping
 -- identity coercions and coalescing nested coercions
-mkCast :: CoreExpr -> Coercion -> CoreExpr
+mkCast :: CoreExpr -> CoercionR -> CoreExpr
 mkCast e co
   | ASSERT2( coercionRole co == Representational
            , text "coercion" <+> ppr co <+> ptext (sLit "passed to mkCast")