Fix map/coerce for newtypes with wrappers (Trac #16208)wip/map-coerce-wrappers

4 files changed, 54 insertions, 0 deletions

diff --git a/compiler/coreSyn/CoreOpt.hs b/compiler/coreSyn/CoreOpt.hs
index 548b5de269..7a8bc66b6e 100644
--- a/compiler/coreSyn/CoreOpt.hs
+++ b/compiler/coreSyn/CoreOpt.hs
@@ -291,6 +291,12 @@ simple_app env (Var v) as
     -- See Note [Unfold compulsory unfoldings in LHSs]
   = simple_app (soeZapSubst env) (unfoldingTemplate unf) as
 
+  | let unf = idUnfolding v
+  , Just a <- isDataConWrapId_maybe v
+  , isNewTyCon (dataConTyCon a)
+    -- See note [Unfold newtype wrappers in LHSs]
+  = simple_app (soeZapSubst env) (unfoldingTemplate unf) as
+
   | otherwise
   , let out_fn = lookupIdSubst (text "simple_app") (soe_subst env) v
   = finish_app env out_fn as
@@ -582,6 +588,19 @@ However, we don't want to inline 'seq', which happens to also have a
 compulsory unfolding, so we only do this unfolding only for things
 that are always-active.  See Note [User-defined RULES for seq] in MkId.
 
+Note [Unfold newtype wrappers in LHSs]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Newtypes may have wrappers, e.g.
+
+newtype Age a b where
+ MkAge :: forall b a. Int -> Age a b
+
+(the wrapper reorders the type arguments)
+
+In order for the `map coerce = coerce` rule to match `map MkAge` (as
+it should!), we need to unfold newtype wrappers in simple_app. See also Note
+[Unfold compulsory unfoldings in LHSs] and Trac #16208.
+
 Note [Getting the map/coerce RULE to work]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We wish to allow the "map/coerce" RULE to fire:
diff --git a/testsuite/tests/simplCore/should_run/T16208.hs b/testsuite/tests/simplCore/should_run/T16208.hs
new file mode 100644
index 0000000000..60c3af7c24
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T16208.hs
@@ -0,0 +1,30 @@
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE RankNTypes #-}
+
+import GHC.Exts
+import Unsafe.Coerce
+
+newtype Age a b where
+  Age :: forall b a. Int -> Age a b
+
+foo :: [Int] -> [Int]
+foo = map id
+fooAge :: [Int] -> [Age a b]
+fooAge = map Age
+fooCoerce :: [Int] -> [Age a b]
+fooCoerce = map coerce
+fooUnsafeCoerce :: [Int] -> [Age a b]
+fooUnsafeCoerce = map unsafeCoerce
+
+same :: a -> b -> IO ()
+same x y = case reallyUnsafePtrEquality# (unsafeCoerce x) y of
+    1# -> putStrLn "yes"
+    _  -> putStrLn "no"
+
+main = do
+    let l = [1,2,3]
+    same (foo l) l
+    same (fooAge l) l
+    same (fooCoerce l) l
+    same (fooUnsafeCoerce l) l
diff --git a/testsuite/tests/simplCore/should_run/T16208.stdout b/testsuite/tests/simplCore/should_run/T16208.stdout
new file mode 100644
index 0000000000..4ff957b404
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T16208.stdout
@@ -0,0 +1,4 @@
+yes
+yes
+yes
+yes
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, [''])