Commit 603b7be7 authored by David Feuer's avatar David Feuer Committed by Herbert Valerio Riedel
Browse files

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:
parent 63899111
......@@ -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:
"amap/coerce" amap coerce = coerce
-- Second functor law:
"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
{-# 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
......@@ -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, [''])
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment