Fix map/coerce rule for newtypes with wrappers

This addresses Trac #16208 by marking newtype wrapper
unfoldings as compulsory.

Furthermore, we can remove the special case for newtypes
in exprIsConApp_maybe (introduced in 7833cf407d1f).

6 files changed, 51 insertions, 29 deletions

diff --git a/compiler/basicTypes/MkId.hs b/compiler/basicTypes/MkId.hs
index ceda50295c..e3b928c4c7 100644
--- a/compiler/basicTypes/MkId.hs
+++ b/compiler/basicTypes/MkId.hs
@@ -298,6 +298,27 @@ so the data constructor for T:C had a single argument, namely the
 predicate (C a).  But now we treat that as an ordinary argument, not
 part of the theta-type, so all is well.
 
+Note [Compulsory newtype unfolding]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Newtype wrappers, just like workers, have compulsory unfoldings.
+This is needed so that two optimizations involving newtypes have the same
+effect whether a wrapper is present or not:
+
+(1) Case-of-known constructor.
+    See Note [beta-reduction in exprIsConApp_maybe].
+
+(2) Matching against the map/coerce RULE. Suppose we have the RULE
+
+    {-# RULE "map/coerce" map coerce = ... #-}
+
+    As described in Note [Getting the map/coerce RULE to work],
+    the occurrence of 'coerce' is transformed into:
+
+    {-# RULE "map/coerce" forall (c :: T1 ~R# T2).
+                          map ((\v -> v) `cast` c) = ... #-}
+
+    We'd like 'map Age' to match the LHS. For this to happen, Age
+    must be unfolded, otherwise we'll be stuck. This is tested in T16208.
 
 ************************************************************************
 *                                                                      *
@@ -607,7 +628,9 @@ mkDataConRep dflags fam_envs wrap_name mb_bangs data_con
              -- See Note [Inline partially-applied constructor wrappers]
              -- Passing Nothing here allows the wrapper to inline when
              -- unsaturated.
-             wrap_unf = mkInlineUnfolding wrap_rhs
+             wrap_unf | isNewTyCon tycon = mkCompulsoryUnfolding wrap_rhs
+                        -- See Note [Compulsory newtype unfolding]
+                      | otherwise        = mkInlineUnfolding wrap_rhs
              wrap_rhs = mkLams wrap_tvs $
                         mkLams wrap_args $
                         wrapFamInstBody tycon res_ty_args $
diff --git a/compiler/coreSyn/CoreOpt.hs b/compiler/coreSyn/CoreOpt.hs
index 80fb3a80cf..d0dba81e3e 100644
--- a/compiler/coreSyn/CoreOpt.hs
+++ b/compiler/coreSyn/CoreOpt.hs
@@ -42,7 +42,7 @@ import OptCoercion ( optCoercion )
 import Type     hiding ( substTy, extendTvSubst, extendCvSubst, extendTvSubstList
                        , isInScope, substTyVarBndr, cloneTyVarBndr )
 import Coercion hiding ( substCo, substCoVarBndr )
-import TyCon        ( tyConArity, isNewTyCon )
+import TyCon        ( tyConArity )
 import TysWiredIn
 import PrelNames
 import BasicTypes
@@ -793,7 +793,7 @@ Here's how exprIsConApp_maybe achieves this:
       scrutinee = (\n. case n of n' -> MkT n') e
 
 2.  Beta-reduce the application, generating a floated 'let'.
-    See Note [Special case for newtype wrappers] below.  Now we have
+    See Note [beta-reduction in exprIsConApp_maybe] below.  Now we have
       scrutinee = case n of n' -> MkT n'
       with floats {Let n = e}
 
@@ -806,8 +806,8 @@ And now we have a known-constructor MkT that we can return.
 Notice that both (2) and (3) require exprIsConApp_maybe to gather and return
 a bunch of floats, both let and case bindings.
 
-Note [Special case for newtype wrappers]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Note [beta-reduction in exprIsConApp_maybe]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 The unfolding a definition (_e.g._ a let-bound variable or a datacon wrapper) is
 typically a function. For instance, take the wrapper for MkT in Note
 [exprIsConApp_maybe on data constructors with wrappers]:
@@ -838,7 +838,8 @@ Is transformed into
 Which, effectively, means emitting a float `let x = arg` and recursively
 analysing the body.
 
-This strategy requires a special case for newtypes. Suppose we have
+For newtypes, this strategy requires that their wrappers have compulsory unfoldings.
+Suppose we have
    newtype T a b where
      MkT :: a -> T b a   -- Note args swapped
 
@@ -853,7 +854,8 @@ This defines a worker function MkT, a wrapper function $WMkT, and an axT:
 
 Now we are optimising
    case $WMkT (I# 3) |> sym axT of I# y -> ...
-we clearly want to simplify this.  The danger is that we'll end up with
+we clearly want to simplify this. If $WMkT did not have a compulsory
+unfolding, we would end up with
    let a = I#3 in case a of I# y -> ...
 because in general, we do this on-the-fly beta-reduction
    (\x. e) blah  -->  let x = blah in e
@@ -863,14 +865,6 @@ But if the case-of-known-constructor doesn't actually fire (i.e.
 exprIsConApp_maybe does not return Just) then nothing happens, and nothing
 will happen the next time either.
 
-For newtype wrappers we know for sure that the argument of the beta-redex
-is used exactly once, so we can substitute aggressively rather than use a let.
-Hence the special case, implemented in dealWithNewtypeWrapper.
-(It's sound for any beta-redex where the argument is used once, of course.)
-
-dealWithNewtypeWrapper is recursive since newtypes can have
-multiple type arguments.
-
 See test T16254, which checks the behavior of newtypes.
 -}
 
@@ -954,12 +948,6 @@ exprIsConApp_maybe (in_scope, id_unf) expr
         = succeedWith in_scope floats $
           pushCoDataCon con args co
 
-        -- See Note [Special case for newtype wrappers]
-        | Just a <- isDataConWrapId_maybe fun
-        , isNewTyCon (dataConTyCon a)
-        , let rhs = uf_tmpl (realIdUnfolding fun)
-        = dealWithNewtypeWrapper (Left in_scope) floats rhs cont
-
         -- Look through data constructor wrappers: they inline late (See Note
         -- [Activation for data constructor wrappers]) but we want to do
         -- case-of-known-constructor optimisation eagerly.
@@ -1005,13 +993,6 @@ exprIsConApp_maybe (in_scope, id_unf) expr
            ; return (in_scope, floats, con, tys, args) }
 
     ----------------------------
-    -- Unconditionally substitute the argument of a newtype
-    dealWithNewtypeWrapper scope floats (Lam v body) (CC (arg:args) co)
-      = dealWithNewtypeWrapper (extend scope v arg) floats body (CC args co)
-    dealWithNewtypeWrapper scope floats expr args
-      = go scope floats expr args
-
-    ----------------------------
     -- Operations on the (Either InScopeSet CoreSubst)
     -- The Left case is wildly dominant
     subst_co (Left {}) co = co
diff --git a/compiler/coreSyn/CoreUtils.hs b/compiler/coreSyn/CoreUtils.hs
index ee79a0f930..5b161995ea 100644
--- a/compiler/coreSyn/CoreUtils.hs
+++ b/compiler/coreSyn/CoreUtils.hs
@@ -1360,7 +1360,6 @@ isExpandableApp fn n_val_args
   | otherwise
   = case idDetails fn of
       DataConWorkId {} -> True  -- Actually handled by isWorkFreeApp
-      DataConWrapId {} -> True  -- See Note [Special case for newtype wrappers]
       RecSelId {}      -> n_val_args == 1  -- See Note [Record selection]
       ClassOpId {}     -> n_val_args == 1
       PrimOpId {}      -> False
diff --git a/testsuite/tests/simplCore/should_run/T16208.hs b/testsuite/tests/simplCore/should_run/T16208.hs
new file mode 100644
index 0000000000..e346ab84f6
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T16208.hs
@@ -0,0 +1,17 @@
+{-# LANGUAGE GADTs, ExplicitForAll #-}
+module Main (main) where
+
+import GHC.Exts
+
+newtype Age a b where
+  Age :: forall b a. Int -> Age a b
+
+data T a = MkT a
+
+{-# NOINLINE foo #-}
+foo :: (Int -> Age Bool Char) -> String
+foo _ = "bad (RULE should have fired)"
+
+{-# RULES "foo/coerce" [1] foo coerce = "good" #-}
+
+main = putStrLn (foo Age)
diff --git a/testsuite/tests/simplCore/should_run/T16208.stdout b/testsuite/tests/simplCore/should_run/T16208.stdout
new file mode 100644
index 0000000000..12799ccbe7
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T16208.stdout
@@ -0,0 +1 @@
+good
diff --git a/testsuite/tests/simplCore/should_run/all.T b/testsuite/tests/simplCore/should_run/all.T
index f8089438c5..646929f778 100644
--- a/testsuite/tests/simplCore/should_run/all.T
+++ b/testsuite/tests/simplCore/should_run/all.T
@@ -50,6 +50,7 @@ test('T5441', [], multimod_compile_and_run, ['T5441', ''])
 test('T5603', reqlib('integer-gmp'), compile_and_run, [''])
 test('T2110', normal, compile_and_run, [''])
 test('AmapCoerce', normal, compile_and_run, [''])
+test('T16208', normal, compile_and_run, [''])
 
 # Run these tests *without* optimisation too
 test('T5625', [ only_ways(['normal','optasm']), exit_code(1) ], compile_and_run, [''])