Fix the implementation of the "push rules"

Richard pointed out (comment:12 of Trac #13025) that my
implementation of the coercion "push rules", newly added
in exprIsConAppMaybe by commit b4c3a66, wasn't quite right.

But in fact that means that the implementation of those same
rules in Simplify.simplCast was wrong too.

Hence this commit:

* Refactor the push rules so they are implemented in just
  one place (CoreSubst.pushCoArgs, pushCoTyArg, pushCoValArg)
  The code in Simplify gets simpler, which is nice.

* Fix the bug that Richard pointed out (to do with hetero-kinded
  coercions)

Then compiler performance worsened, which led mt do discover
two performance bugs:

* The smart constructor Coercion.mkNthCo didn't have a case
  for ForAllCos, which meant we stupidly build a complicated
  coercion where a simple one would do

* In OptCoercion there was one place where we used CoherenceCo
  (the data constructor) rather than mkCoherenceCo (the smart
  constructor), which meant that the the stupid complicated
  coercion wasn't optimised away

For reasons I don't fully understand, T5321Fun did 2% less compiler
allocation after all this, which is good.

5 files changed, 202 insertions, 153 deletions

diff --git a/compiler/coreSyn/CoreSubst.hs b/compiler/coreSyn/CoreSubst.hs
index e4f2f59a5a..73be490edb 100644
--- a/compiler/coreSyn/CoreSubst.hs
+++ b/compiler/coreSyn/CoreSubst.hs
@@ -34,6 +34,7 @@ module CoreSubst (
         -- ** Simple expression optimiser
         simpleOptPgm, simpleOptExpr, simpleOptExprWith,
         exprIsConApp_maybe, exprIsLiteral_maybe, exprIsLambda_maybe,
+        pushCoArg, pushCoValArg, pushCoTyArg
     ) where
 
 #include "HsVersions.h"
@@ -1245,13 +1246,13 @@ exprIsConApp_maybe (in_scope, id_unf) expr
 
         | Just con <- isDataConWorkId_maybe fun
         , count isValArg args == idArity fun
-        = dealWithCoercion co con args
+        = pushCoDataCon con args co
 
         -- Look through dictionary functions; see Note [Unfolding DFuns]
         | DFunUnfolding { df_bndrs = bndrs, df_con = con, df_args = dfun_args } <- unfolding
         , bndrs `equalLength` args    -- See Note [DFun arity check]
         , let subst = mkOpenSubst in_scope (bndrs `zip` args)
-        = dealWithCoercion co con (map (substExpr (text "exprIsConApp1") subst) dfun_args)
+        = pushCoDataCon con (map (substExpr (text "exprIsConApp1") subst) dfun_args) co
 
         -- Look through unfoldings, but only arity-zero one;
         -- if arity > 0 we are effectively inlining a function call,
@@ -1284,35 +1285,6 @@ exprIsConApp_maybe (in_scope, id_unf) expr
     extend (Left in_scope) v e = Right (extendSubst (mkEmptySubst in_scope) v e)
     extend (Right s)       v e = Right (extendSubst s v e)
 
-pushCoArgs :: Coercion -> [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)
--- 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
--- C.f. simplCast in Simplify.hs
-
-pushCoArg co arg
-  = case arg of
-      Type ty | isForAllTy tyL
-        -> ASSERT2( isForAllTy tyR, ppr co $$ ppr ty )
-           Just (Type ty, mkInstCo co (mkNomReflCo ty))
-
-      _ | isFunTy tyL
-        , [co1, co2] <- decomposeCo 2 co
-              -- If   co  :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
-              -- then co1 :: tyL1 ~ tyR1
-              --      co2 :: tyL2 ~ tyR2
-        -> ASSERT2( isFunTy tyR, ppr co $$ ppr arg )
-           Just (mkCast arg (mkSymCo co1), co2)
-
-      _ -> Nothing
-  where
-    Pair tyL tyR = coercionKind co
 
 -- See Note [exprIsConApp_maybe on literal strings]
 dealWithStringLiteral :: Var -> BS.ByteString -> Coercion
@@ -1323,7 +1295,7 @@ dealWithStringLiteral :: Var -> BS.ByteString -> Coercion
 -- generates a string literal directly.
 dealWithStringLiteral _   str co
   | BS.null str
-  = dealWithCoercion co nilDataCon [Type charTy]
+  = pushCoDataCon nilDataCon [Type charTy] co
 
 dealWithStringLiteral fun str co
   = let strFS = mkFastStringByteString str
@@ -1337,67 +1309,7 @@ dealWithStringLiteral fun str co
                  else App (Var fun)
                           (Lit (MachStr charTail))
 
-    in dealWithCoercion co consDataCon [Type charTy, char, rest]
-
-dealWithCoercion :: Coercion -> DataCon -> [CoreExpr]
-                 -> Maybe (DataCon, [Type], [CoreExpr])
-dealWithCoercion co dc dc_args
-  | isReflCo co || from_ty `eqType` to_ty  -- try cheap test first
-  , let (univ_ty_args, rest_args) = splitAtList (dataConUnivTyVars dc) dc_args
-  = Just (dc, map exprToType univ_ty_args, rest_args)
-
-  | Just (to_tc, to_tc_arg_tys) <- splitTyConApp_maybe to_ty
-  , to_tc == dataConTyCon dc
-        -- These two tests can fail; we might see
-        --      (C x y) `cast` (g :: T a ~ S [a]),
-        -- where S is a type function.  In fact, exprIsConApp
-        -- will probably not be called in such circumstances,
-        -- but there't nothing wrong with it
-
-  =     -- Here we do the KPush reduction rule as described in "Down with kinds"
-        -- The transformation applies iff we have
-        --      (C e1 ... en) `cast` co
-        -- where co :: (T t1 .. tn) ~ to_ty
-        -- The left-hand one must be a T, because exprIsConApp returned True
-        -- but the right-hand one might not be.  (Though it usually will.)
-    let
-        tc_arity       = tyConArity to_tc
-        dc_univ_tyvars = dataConUnivTyVars dc
-        dc_ex_tyvars   = dataConExTyVars dc
-        arg_tys        = dataConRepArgTys dc
-
-        non_univ_args  = dropList dc_univ_tyvars dc_args
-        (ex_args, val_args) = splitAtList dc_ex_tyvars non_univ_args
-
-        -- Make the "Psi" from the paper
-        omegas = decomposeCo tc_arity co
-        (psi_subst, to_ex_arg_tys)
-          = liftCoSubstWithEx Representational
-                              dc_univ_tyvars
-                              omegas
-                              dc_ex_tyvars
-                              (map exprToType ex_args)
-
-          -- Cast the value arguments (which include dictionaries)
-        new_val_args = zipWith cast_arg arg_tys val_args
-        cast_arg arg_ty arg = mkCast arg (psi_subst arg_ty)
-
-        to_ex_args = map Type to_ex_arg_tys
-
-        dump_doc = vcat [ppr dc,      ppr dc_univ_tyvars, ppr dc_ex_tyvars,
-                         ppr arg_tys, ppr dc_args,
-                         ppr ex_args, ppr val_args, ppr co, ppr from_ty, ppr to_ty, ppr to_tc ]
-    in
-    ASSERT2( eqType from_ty (mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args)), dump_doc )
-    ASSERT2( equalLength val_args arg_tys, dump_doc )
-    Just (dc, to_tc_arg_tys, to_ex_args ++ new_val_args)
-
-  | otherwise
-  = Nothing
-
-  where
-    Pair from_ty to_ty = coercionKind co
-
+    in pushCoDataCon consDataCon [Type charTy, char, rest] co
 
 {-
 Note [Unfolding DFuns]
@@ -1489,11 +1401,107 @@ exprIsLambda_maybe _ _e
       Nothing
 
 
+{- *********************************************************************
+*                                                                      *
+              The "push rules"
+*                                                                      *
+************************************************************************
+
+Here we implement the "push rules" from FC papers:
+
+* The push-argument ules, where we can move a coercion past an argument.
+  We have
+      (fun |> co) arg
+  and we want to transform it to
+    (fun arg') |> co'
+  for some suitable co' and tranformed arg'.
+
+* The PushK rule for data constructors.  We have
+       (K e1 .. en) |> co
+  and we want to tranform to
+       (K e1' .. en')
+  by pushing the coercion into the oarguments
+-}
+
+pushCoArgs :: Coercion -> [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)
+-- 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
+-- C.f. simplCast in Simplify.hs
+-- 'co' is always Representational
+
+pushCoArg co (Type ty) = do { (ty', co') <- pushCoTyArg co ty
+                            ; return (Type ty', co') }
+pushCoArg co val_arg   = do { (arg_co, co') <- pushCoValArg co
+                            ; return (mkCast val_arg arg_co, co') }
+
+pushCoTyArg :: Coercion -> Type -> Maybe (Type, Coercion)
+-- We have (fun |> co) @ty
+-- Push the coercion through to return
+--         (fun @ty') |> co'
+-- 'co' is always Representational
+pushCoTyArg co ty
+  | tyL `eqType` tyR
+  = Just (ty, mkRepReflCo (piResultTy tyR ty))
+
+  | isForAllTy tyL
+  = ASSERT2( isForAllTy tyR, ppr co $$ ppr ty )
+    Just (ty `mkCastTy` mkSymCo co1, co2)
+
+  | otherwise
+  = Nothing
+  where
+    Pair tyL tyR = coercionKind co
+       -- co :: tyL ~ tyR
+       -- tyL = forall (a1 :: k1). ty1
+       -- tyR = forall (a2 :: k2). ty2
+
+    co1 = mkNthCo 0 co
+       -- co1 :: k1 ~ k2
+       -- Note that NthCo can extract an equality between the kinds
+       -- of the types related by a coercion between forall-types.
+       -- See the NthCo case in CoreLint.
+
+    co2 = mkInstCo co (mkCoherenceLeftCo (mkNomReflCo ty) co1)
+        -- co2 :: ty1[ (ty|>co1)/a1 ] ~ ty2[ ty/a2 ]
+        -- Arg of mkInstCo is always nominal, hence mkNomReflCo
+
+pushCoValArg :: Coercion -> Maybe (Coercion, Coercion)
+-- We have (fun |> co) arg
+-- Push the coercion through to return
+--         (fun (arg |> co_arg)) |> co_res
+-- 'co' is always Representational
+pushCoValArg co
+  | tyL `eqType` tyR
+  = Just (mkRepReflCo arg, mkRepReflCo res)
+
+  | isFunTy tyL
+  , [co1, co2] <- decomposeCo 2 co
+              -- If   co  :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
+              -- then co1 :: tyL1 ~ tyR1
+              --      co2 :: tyL2 ~ tyR2
+  = ASSERT2( isFunTy tyR, ppr co $$ ppr arg )
+    Just (mkSymCo co1, co2)
+
+  | otherwise
+  = Nothing
+  where
+    (arg, res)   = splitFunTy tyR
+    Pair tyL tyR = coercionKind co
+
 pushCoercionIntoLambda
     :: InScopeSet -> Var -> CoreExpr -> Coercion -> Maybe (Var, CoreExpr)
+-- This implements the Push rule from the paper on coercions
+--    (\x. e) |> co
+-- ===>
+--    (\x'. e |> co')
 pushCoercionIntoLambda in_scope x e co
-    -- This implements the Push rule from the paper on coercions
-    -- Compare with simplCast in Simplify
     | ASSERT(not (isTyVar x) && not (isCoVar x)) True
     , Pair s1s2 t1t2 <- coercionKind co
     , Just (_s1,_s2) <- splitFunTy_maybe s1s2
@@ -1510,3 +1518,64 @@ pushCoercionIntoLambda in_scope x e co
     | otherwise
     = pprTrace "exprIsLambda_maybe: Unexpected lambda in case" (ppr (Lam x e))
       Nothing
+
+pushCoDataCon :: DataCon -> [CoreExpr] -> Coercion
+              -> Maybe (DataCon
+                       , [Type]      -- Universal type args
+                       , [CoreExpr]) -- All other args incl existentials
+-- Implement the KPush reduction rule as described in "Down with kinds"
+-- The transformation applies iff we have
+--      (C e1 ... en) `cast` co
+-- where co :: (T t1 .. tn) ~ to_ty
+-- The left-hand one must be a T, because exprIsConApp returned True
+-- but the right-hand one might not be.  (Though it usually will.)
+pushCoDataCon dc dc_args co
+  | isReflCo co || from_ty `eqType` to_ty  -- try cheap test first
+  , let (univ_ty_args, rest_args) = splitAtList (dataConUnivTyVars dc) dc_args
+  = Just (dc, map exprToType univ_ty_args, rest_args)
+
+  | Just (to_tc, to_tc_arg_tys) <- splitTyConApp_maybe to_ty
+  , to_tc == dataConTyCon dc
+        -- These two tests can fail; we might see
+        --      (C x y) `cast` (g :: T a ~ S [a]),
+        -- where S is a type function.  In fact, exprIsConApp
+        -- will probably not be called in such circumstances,
+        -- but there't nothing wrong with it
+
+  = let
+        tc_arity       = tyConArity to_tc
+        dc_univ_tyvars = dataConUnivTyVars dc
+        dc_ex_tyvars   = dataConExTyVars dc
+        arg_tys        = dataConRepArgTys dc
+
+        non_univ_args  = dropList dc_univ_tyvars dc_args
+        (ex_args, val_args) = splitAtList dc_ex_tyvars non_univ_args
+
+        -- Make the "Psi" from the paper
+        omegas = decomposeCo tc_arity co
+        (psi_subst, to_ex_arg_tys)
+          = liftCoSubstWithEx Representational
+                              dc_univ_tyvars
+                              omegas
+                              dc_ex_tyvars
+                              (map exprToType ex_args)
+
+          -- Cast the value arguments (which include dictionaries)
+        new_val_args = zipWith cast_arg arg_tys val_args
+        cast_arg arg_ty arg = mkCast arg (psi_subst arg_ty)
+
+        to_ex_args = map Type to_ex_arg_tys
+
+        dump_doc = vcat [ppr dc,      ppr dc_univ_tyvars, ppr dc_ex_tyvars,
+                         ppr arg_tys, ppr dc_args,
+                         ppr ex_args, ppr val_args, ppr co, ppr from_ty, ppr to_ty, ppr to_tc ]
+    in
+    ASSERT2( eqType from_ty (mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args)), dump_doc )
+    ASSERT2( equalLength val_args arg_tys, dump_doc )
+    Just (dc, to_tc_arg_tys, to_ex_args ++ new_val_args)
+
+  | otherwise
+  = Nothing
+
+  where
+    Pair from_ty to_ty = coercionKind co
diff --git a/compiler/simplCore/Simplify.hs b/compiler/simplCore/Simplify.hs
index 0b9f5f0948..aaeb997b54 100644
--- a/compiler/simplCore/Simplify.hs
+++ b/compiler/simplCore/Simplify.hs
@@ -35,6 +35,7 @@ import PprCore          ( pprCoreExpr )
 import CoreUnfold
 import CoreUtils
 import CoreArity
+import CoreSubst        ( pushCoTyArg, pushCoValArg )
 --import PrimOp           ( tagToEnumKey ) -- temporalily commented out. See #8326
 import Rules            ( mkRuleInfo, lookupRule, getRules )
 import TysPrim          ( voidPrimTy ) --, intPrimTy ) -- temporalily commented out. See #8326
@@ -1144,65 +1145,38 @@ simplCast env body co0 cont0
         ; cont1 <- addCoerce co1 cont0
         ; simplExprF env body cont1 }
   where
-       addCoerce co cont = add_coerce co (coercionKind co) cont
-
-       add_coerce _co (Pair s1 k1) cont     -- co :: ty~ty
-         | s1 `eqType` k1 = return cont    -- is a no-op
-
-       add_coerce co1 (Pair s1 _k2) (CastIt co2 cont)
-         | (Pair _l1 t1) <- coercionKind co2
-                --      e |> (g1 :: S1~L) |> (g2 :: L~T1)
-                -- ==>
-                --      e,                       if S1=T1
-                --      e |> (g1 . g2 :: S1~T1)  otherwise
-                --
+       addCoerce :: OutCoercion -> SimplCont -> SimplM SimplCont
+       addCoerce co1 (CastIt co2 cont)
+         = addCoerce (mkTransCo co1 co2) cont
+
+       addCoerce co cont@(ApplyToTy { sc_arg_ty = arg_ty, sc_cont = tail })
+         | Just (arg_ty', co') <- pushCoTyArg co arg_ty
+         = do { tail' <- addCoerce co' tail
+              ; return (cont { sc_arg_ty = arg_ty', sc_cont = tail' }) }
+
+       addCoerce co (ApplyToVal { sc_arg = arg, sc_env = arg_se
+                                , sc_dup = dup, sc_cont = tail })
+         | Just (co1, co2) <- pushCoValArg co
+         = do { (dup', arg_se', arg') <- simplArg env dup arg_se arg
+                   -- When we build the ApplyTo we can't mix the OutCoercion
+                   -- 'co' with the InExpr 'arg', so we simplify
+                   -- to make it all consistent.  It's a bit messy.
+                   -- But it isn't a common case.
+                   -- Example of use: Trac #995
+              ; tail' <- addCoerce co2 tail
+              ; return (ApplyToVal { sc_arg  = mkCast arg' co1
+                                   , sc_env  = arg_se'
+                                   , sc_dup  = dup'
+                                   , sc_cont = tail' }) }
+
+       addCoerce co cont
+         | isReflexiveCo co = return cont
+         | otherwise        = return (CastIt co cont)
+                -- It's worth checking isReflexiveCo.
                 -- For example, in the initial form of a worker
                 -- we may find  (coerce T (coerce S (\x.e))) y
                 -- and we'd like it to simplify to e[y/x] in one round
                 -- of simplification
-         , s1 `eqType` t1  = return cont     -- The coerces cancel out
-         | otherwise       = return (CastIt (mkTransCo co1 co2) cont)
-
-       add_coerce co (Pair s1s2 _t1t2) cont@(ApplyToTy { sc_arg_ty = arg_ty, sc_cont = tail })
-                -- (f |> g) ty  --->   (f ty) |> (g @ ty)
-                -- This implements the PushT rule from the paper
-         | isForAllTy s1s2
-         = do { cont' <- addCoerce new_cast tail
-              ; return (cont { sc_cont = cont' }) }
-         where
-           new_cast = mkInstCo co (mkNomReflCo arg_ty)
-
-       add_coerce co (Pair s1s2 t1t2) (ApplyToVal { sc_arg = arg, sc_env = arg_se
-                                                  , sc_dup = dup, sc_cont = cont })
-         | isFunTy s1s2   -- This implements the Push rule from the paper
-         , isFunTy t1t2   -- Check t1t2 to ensure 'arg' is a value arg
-                --      (e |> (g :: s1s2 ~ t1->t2)) f
-                -- ===>
-                --      (e (f |> (arg g :: t1~s1))
-                --      |> (res g :: s2->t2)
-                --
-                -- t1t2 must be a function type, t1->t2, because it's applied
-                -- to something but s1s2 might conceivably not be
-                --
-                -- When we build the ApplyTo we can't mix the out-types
-                -- with the InExpr in the argument, so we simply substitute
-                -- to make it all consistent.  It's a bit messy.
-                -- But it isn't a common case.
-                --
-                -- Example of use: Trac #995
-         = do { (dup', arg_se', arg') <- simplArg env dup arg_se arg
-              ; cont'                <- addCoerce co2 cont
-              ; return (ApplyToVal { sc_arg  = mkCast arg' (mkSymCo co1)
-                                   , sc_env  = arg_se'
-                                   , sc_dup  = dup'
-                                   , sc_cont = cont' }) }
-         where
-           -- we split coercion t1->t2 ~ s1->s2 into t1 ~ s1 and
-           -- t2 ~ s2 with left and right on the curried form:
-           --    (->) t1 t2 ~ (->) s1 s2
-           [co1, co2] = decomposeCo 2 co
-
-       add_coerce co _ cont = return (CastIt co cont)
 
 simplArg :: SimplEnv -> DupFlag -> StaticEnv -> CoreExpr
          -> SimplM (DupFlag, StaticEnv, OutExpr)
diff --git a/compiler/types/Coercion.hs b/compiler/types/Coercion.hs
index 0adadc33ca..8b6eccb6eb 100644
--- a/compiler/types/Coercion.hs
+++ b/compiler/types/Coercion.hs
@@ -853,7 +853,10 @@ mkNthCo n (Refl r ty)
           | otherwise
           = False
 
-mkNthCo n (TyConAppCo _ _ cos) = cos `getNth` n
+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 = NthCo n co
 
 mkLRCo :: LeftOrRight -> Coercion -> Coercion
diff --git a/compiler/types/OptCoercion.hs b/compiler/types/OptCoercion.hs
index ca67bc7274..a7dadf39e0 100644
--- a/compiler/types/OptCoercion.hs
+++ b/compiler/types/OptCoercion.hs
@@ -301,7 +301,7 @@ opt_co4 env sym rep r (CoherenceCo co1 co2)
            else opt_trans in_scope (mkCoherenceCo col1' co2') cor1'
 
   | otherwise
-  = wrapSym sym $ CoherenceCo (opt_co4_wrap env False rep r co1) co2'
+  = wrapSym sym $ mkCoherenceCo (opt_co4_wrap env False rep r co1) co2'
   where co1' = opt_co4_wrap env sym   rep   r       co1
         co2' = opt_co4_wrap env False False Nominal co2
         in_scope = lcInScopeSet env
diff --git a/testsuite/tests/perf/compiler/all.T b/testsuite/tests/perf/compiler/all.T
index 87a6c8184d..c9ed905926 100644
--- a/testsuite/tests/perf/compiler/all.T
+++ b/testsuite/tests/perf/compiler/all.T
@@ -501,7 +501,7 @@ test('T5321Fun',
              # 2014-09-03: 299656164     (specialisation and inlining)
              # 10/12/2014: 206406188     #  Improvements in constraint solver
              # 2016-04-06: 279922360 x86/Linux
-            (wordsize(64), 565883176, 10)])
+            (wordsize(64), 497356688, 5)])
              # prev:       585521080
              # 29/08/2012: 713385808     #  (increase due to new codegen)
              # 15/05/2013: 628341952     #  (reason for decrease unknown)
@@ -515,6 +515,9 @@ test('T5321Fun',
              #  (undefined now takes an implicit parameter and GHC -O0 does
              #  not recognize that the application is bottom)
              # 11/12/2015: 565883176     #  TypeInType (see #11196)
+             # 06/01/2017: 497356688     #  Small coercion optimisations
+                                         #  The actual decrease was only 2%; earlier
+                                         #    commits had drifted down
       ],
       compile,[''])