Push coercions in exprIsConApp_maybe

Trac #13025 showed up the fact that exprIsConApp_maybe isn't
clever enough: it didn't push coercions through applicatins, and that
meant we weren't getting as much superclass selection as we should.

It's easy to fix, happily.

See Note [Push coercions in exprIsConApp_maybe]

6 files changed, 114 insertions, 2 deletions

diff --git a/compiler/coreSyn/CoreSubst.hs b/compiler/coreSyn/CoreSubst.hs
index e8a8f6ea55..e4f2f59a5a 100644
--- a/compiler/coreSyn/CoreSubst.hs
+++ b/compiler/coreSyn/CoreSubst.hs
@@ -1196,6 +1196,18 @@ Just (':', [Char], ['a', unpackCString# "bc"]).
 We need to be careful about UTF8 strings here. ""# contains a ByteString, so
 we must parse it back into a FastString to split off the first character.
 That way we can treat unpackCString# and unpackCStringUtf8# in the same way.
+
+Note [Push coercions in exprIsConApp_maybe]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+In Trac #13025 I found a case where we had
+    op (df @t1 @t2)     -- op is a ClassOp
+where
+    df = (/\a b. K e1 e2) |> g
+
+To get this to come out we need to simplify on the fly
+   ((/\a b. K e1 e2) |> g) @t1 @t2
+
+Hence the use of pushCoArgs.
 -}
 
 data ConCont = CC [CoreExpr] Coercion
@@ -1209,12 +1221,16 @@ exprIsConApp_maybe (in_scope, id_unf) expr
   = go (Left in_scope) expr (CC [] (mkRepReflCo (exprType expr)))
   where
     go :: Either InScopeSet Subst
+             -- Left in-scope  means "empty substitution"
+             -- Right subst    means "apply this substitution to the CoreExpr"
        -> CoreExpr -> ConCont
        -> Maybe (DataCon, [Type], [CoreExpr])
     go subst (Tick t expr) cont
        | not (tickishIsCode t) = go subst expr cont
-    go subst (Cast expr co1) (CC [] co2)
-       = go subst expr (CC [] (subst_co subst co1 `mkTransCo` co2))
+    go subst (Cast expr co1) (CC args co2)
+       | Just (args', co1') <- pushCoArgs (subst_co subst co1) args
+            -- See Note [Push coercions in exprIsConApp_maybe]
+       = go subst expr (CC args' (co1' `mkTransCo` co2))
     go subst (App fun arg) (CC args co)
        = go subst fun (CC (subst_arg subst arg : args) co)
     go subst (Lam var body) (CC (arg:args) co)
@@ -1268,6 +1284,36 @@ 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
                       -> Maybe (DataCon, [Type], [CoreExpr])
diff --git a/testsuite/tests/simplCore/should_compile/Makefile b/testsuite/tests/simplCore/should_compile/Makefile
index a5d9a1e34a..5791daff00 100644
--- a/testsuite/tests/simplCore/should_compile/Makefile
+++ b/testsuite/tests/simplCore/should_compile/Makefile
@@ -166,3 +166,9 @@ T5615:
 	-grep 'quotInt#' T5615.dump-simpl
 	-grep 'remInt#' T5615.dump-simpl
 	grep -c '1999#' T5615.dump-simpl
+
+T13025:
+	$(RM) -f T13025.o T13025.hi T13025a.o T13025a.hi
+	'$(TEST_HC)' $(TEST_HC_OPTS) -c -O T13025a.hs
+	'$(TEST_HC)' $(TEST_HC_OPTS) -c -O T13025.hs -ddump-simpl | grep HEq_sc | wc
+	# No lines should match 'HEq_sc' so wc should output zeros
diff --git a/testsuite/tests/simplCore/should_compile/T13025.hs b/testsuite/tests/simplCore/should_compile/T13025.hs
new file mode 100644
index 0000000000..01facb8601
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T13025.hs
@@ -0,0 +1,15 @@
+{-# LANGUAGE DataKinds #-}
+module T13025 where
+import T13025a
+
+type MyRec = Rec '[ '("A",Int), '("B",Int), '("C",Int) ]
+
+getC :: MyRec -> Int
+getC = getField (Proxy::Proxy '("C",Int))
+
+doubleC :: MyRec -> MyRec
+doubleC r = setC (2 * (getC r)) r
+  where setC = set . (Field :: Int -> Field '("C",Int))
+
+main :: IO ()
+main = print (getC (Field 1 :& Field 2 :& Field 3 :& Nil :: MyRec))
diff --git a/testsuite/tests/simplCore/should_compile/T13025.stdout b/testsuite/tests/simplCore/should_compile/T13025.stdout
new file mode 100644
index 0000000000..7d1413f1c1
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T13025.stdout
@@ -0,0 +1 @@
+      0       0       0
diff --git a/testsuite/tests/simplCore/should_compile/T13025a.hs b/testsuite/tests/simplCore/should_compile/T13025a.hs
new file mode 100644
index 0000000000..3f9a4cbe21
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T13025a.hs
@@ -0,0 +1,40 @@
+{-# LANGUAGE ConstraintKinds, DataKinds, FlexibleContexts,
+             FlexibleInstances, GADTs, MultiParamTypeClasses,
+             PolyKinds, ScopedTypeVariables, TypeFamilies,
+             TypeOperators #-}
+module T13025a where
+
+data Nat = Z | S Nat
+data Proxy a = Proxy
+
+data Field :: (k,*) -> * where
+  Field :: a -> Field '(s,a)
+
+type family Index r rs :: Nat where
+  Index r (r ': rs) = 'Z
+  Index r (s ': rs) = 'S (Index r rs)
+
+data Rec (rs :: [ (k,*) ]) where
+  Nil :: Rec '[]
+  (:&) :: Field r -> Rec rs -> Rec (r ': rs)
+infixr 5 :&
+
+class Index r rs ~ i => HasField r rs i where
+  get :: proxy r -> Rec rs -> Field r
+  set :: Field r -> Rec rs -> Rec rs
+
+instance HasField r (r ': rs) 'Z where
+  get _ (x :& _) = x
+  set x (_ :& xs) = x :& xs
+
+instance (HasField r rs i, Index r (s ': rs) ~ 'S i)
+         => HasField r (s ': rs) ('S i) where
+  get p (_ :& xs) = get p xs
+  set x' (x :& xs) = x :& set x' xs
+
+type Has r rs = HasField r rs (Index r rs)
+
+getField :: Has '(s,a) rs => proxy '(s,a) -> Rec rs -> a
+getField p = aux . get p
+  where aux :: Field '(s,a) -> a
+        aux (Field x) = x
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index c5666c4a6c..e09880fb67 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -255,4 +255,8 @@ test('T12603',
      run_command,
      ['$MAKE -s --no-print-directory T12603'])
 test('T13027', normal, compile, [''])
+test('T13025',
+     normal,
+     run_command,
+     ['$MAKE -s --no-print-directory T13025'])