Implement amap/coerce for Array (re #9796)

Implement an `amap`/`coerce` rule in `GHC.Arr` to match the
`map`/`coerce` rule in GHC.Base.

In order to do so, delay inlining `amap` until phase 1.

To prevent the inlining delay from causing major inefficiencies due to
missed list fusion, rewrite `amap` to avoid relying on list fusion. This
has the extra benefit of reducing the size of the compiled amap code by
skipping the impossible case of an array with a negative size.

Reviewed By: nomeata

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

4 files changed, 66 insertions, 3 deletions

diff --git a/libraries/base/GHC/Arr.hs b/libraries/base/GHC/Arr.hs
index 67702ea884..02bf7d8a61 100644
--- a/libraries/base/GHC/Arr.hs
+++ b/libraries/base/GHC/Arr.hs
@@ -704,10 +704,44 @@ unsafeAccum f arr ies = runST (do
     STArray l u n marr# <- thawSTArray arr
     ST (foldr (adjust f marr#) (done l u n marr#) ies))
 
-{-# INLINE amap #-}
+{-# INLINE [1] amap #-}
 amap :: Ix i => (a -> b) -> Array i a -> Array i b
-amap f arr@(Array l u n _) =
-    unsafeArray' (l,u) n [(i, f (unsafeAt arr i)) | i <- [0 .. n - 1]]
+amap f arr@(Array l u n@(I# n#) _) = runST (ST $ \s1# ->
+    case newArray# n# arrEleBottom s1# of
+        (# s2#, marr# #) ->
+          let go i s#
+                | i == n    = done l u n marr# s#
+                | otherwise = fill marr# (i, f (unsafeAt arr i)) (go (i+1)) s#
+          in go 0 s2# )
+
+{-
+amap was originally defined like this:
+
+ amap f arr@(Array l u n _) =
+     unsafeArray' (l,u) n [(i, f (unsafeAt arr i)) | i <- [0 .. n - 1]]
+
+There are two problems:
+
+1. The enumFromTo implementation produces (spurious) code for the impossible
+case of n<0 that ends up duplicating the array freezing code.
+
+2. This implementation relies on list fusion for efficiency. In order to
+implement the amap/coerce rule, we need to delay inlining amap until simplifier
+phase 1, which is when the eftIntList rule kicks in and makes that impossible.
+-}
+
+
+-- See Breitner, Eisenberg, Peyton Jones, and Weirich, "Safe Zero-cost
+-- Coercions for Haskell", section 6.5:
+--   http://research.microsoft.com/en-us/um/people/simonpj/papers/ext-f/coercible.pdf
+{-# RULES
+"amap/coerce" amap coerce = coerce
+ #-}
+
+-- Second functor law:
+{-# RULES
+"amap/amap" forall f g a . amap f (amap g a) = amap (f . g) a
+ #-}
 
 -- | 'ixmap' allows for transformations on array indices.
 -- It may be thought of as providing function composition on the right
diff --git a/testsuite/tests/simplCore/should_run/AmapCoerce.hs b/testsuite/tests/simplCore/should_run/AmapCoerce.hs
new file mode 100644
index 0000000000..01a9a5d5c6
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/AmapCoerce.hs
@@ -0,0 +1,25 @@
+{-# LANGUAGE MagicHash #-}
+
+import GHC.Exts
+import Unsafe.Coerce
+import Data.Array
+
+newtype Age = Age Int
+
+fooAge :: Array Int Int -> Array Int Age
+fooAge = fmap Age
+fooCoerce :: Array Int Int -> Array Int Age
+fooCoerce = fmap coerce
+fooUnsafeCoerce :: Array Int Int -> Array Int Age
+fooUnsafeCoerce = fmap unsafeCoerce
+
+same :: a -> b -> IO ()
+same x y = case reallyUnsafePtrEquality# (unsafeCoerce x) y of
+    1# -> putStrLn "yes"
+    _  -> putStrLn "no"
+
+main = do
+    let l = listArray (1,3) [1,2,3]
+    same (fooAge l) l
+    same (fooCoerce l) l
+    same (fooUnsafeCoerce l) l
diff --git a/testsuite/tests/simplCore/should_run/AmapCoerce.stdout b/testsuite/tests/simplCore/should_run/AmapCoerce.stdout
new file mode 100644
index 0000000000..55f7ebb441
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/AmapCoerce.stdout
@@ -0,0 +1,3 @@
+yes
+yes
+yes
diff --git a/testsuite/tests/simplCore/should_run/all.T b/testsuite/tests/simplCore/should_run/all.T
index 93dc4c66f9..364dfd694f 100644
--- a/testsuite/tests/simplCore/should_run/all.T
+++ b/testsuite/tests/simplCore/should_run/all.T
@@ -53,6 +53,7 @@ test('T5441', extra_clean(['T5441a.o','T5441a.hi']),
               multimod_compile_and_run, ['T5441',''])
 test('T5603', normal, compile_and_run, [''])
 test('T2110', normal, compile_and_run, [''])
+test('AmapCoerce', normal, compile_and_run, [''])
 
 # Run these tests *without* optimisation too
 test('T5625', [ only_ways(['normal','optasm']), exit_code(1) ], compile_and_run, [''])