[commit: ghc] master: Make accumArray and accum stricter (08345bd)

[git at git.haskell.org]

Repository : ssh://git@git.haskell.org/ghc

On branch  : master
Link       : http://ghc.haskell.org/trac/ghc/changeset/08345bd0e8d237ec3929aaee7613c4f76e07e131/ghc

>---------------------------------------------------------------

commit 08345bd0e8d237ec3929aaee7613c4f76e07e131
Author: David Feuer <david.feuer at gmail.com>
Date:   Mon Mar 5 15:18:05 2018 -0500

    Make accumArray and accum stricter
    
    `accumArray` was lazier than documented. `accum` did not have
    documented strictness. The extra laziness allowed thunks to build
    up in the array. Force the results of applying the accumulating
    function to resolve.
    
    Reviewers: hvr, bgamari
    
    Reviewed By: bgamari
    
    Subscribers: alpmestan, rwbarton, thomie, carter
    
    GHC Trac Issues: #14785
    
    Differential Revision: https://phabricator.haskell.org/D4403


>---------------------------------------------------------------

08345bd0e8d237ec3929aaee7613c4f76e07e131
 libraries/base/GHC/Arr.hs | 30 +++++++++++++++++++++++-------
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/libraries/base/GHC/Arr.hs b/libraries/base/GHC/Arr.hs
index adfd602..3698852 100644
--- a/libraries/base/GHC/Arr.hs
+++ b/libraries/base/GHC/Arr.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE Unsafe #-}
 {-# LANGUAGE NoImplicitPrelude, MagicHash, UnboxedTuples, RoleAnnotations #-}
+{-# LANGUAGE BangPatterns #-}
 {-# OPTIONS_HADDOCK hide #-}
 
 -----------------------------------------------------------------------------
@@ -505,7 +506,7 @@ listArray (l,u) es = runST (ST $ \s1# ->
 -- | The value at the given index in an array.
 {-# INLINE (!) #-}
 (!) :: Ix i => Array i e -> i -> e
-arr@(Array l u n _) ! i = unsafeAt arr $ safeIndex (l,u) n i
+(!) arr@(Array l u n _) i = unsafeAt arr $ safeIndex (l,u) n i
 
 {-# INLINE safeRangeSize #-}
 safeRangeSize :: Ix i => (i, i) -> Int
@@ -636,6 +637,7 @@ assocs arr@(Array l u _ _) =
 -- | The 'accumArray' function deals with repeated indices in the association
 -- list using an /accumulating function/ which combines the values of
 -- associations with the same index.
+--
 -- For example, given a list of values of some index type, @hist@
 -- produces a histogram of the number of occurrences of each index within
 -- a specified range:
@@ -643,10 +645,10 @@ assocs arr@(Array l u _ _) =
 -- > hist :: (Ix a, Num b) => (a,a) -> [a] -> Array a b
 -- > hist bnds is = accumArray (+) 0 bnds [(i, 1) | i<-is, inRange bnds i]
 --
--- If the accumulating function is strict, then 'accumArray' is strict in
--- the values, as well as the indices, in the association list.  Thus,
--- unlike ordinary arrays built with 'array', accumulated arrays should
--- not in general be recursive.
+-- @accumArray@ is strict in each result of applying the accumulating
+-- function, although it is lazy in the initial value. Thus, unlike
+-- arrays built with 'array', accumulated arrays should not in general
+-- be recursive.
 {-# INLINE accumArray #-}
 accumArray :: Ix i
         => (e -> a -> e)        -- ^ accumulating function
@@ -667,7 +669,7 @@ unsafeAccumArray f initial b ies = unsafeAccumArray' f initial b (rangeSize b) i
 unsafeAccumArray' :: (e -> a -> e) -> e -> (i,i) -> Int -> [(Int, a)] -> Array i e
 unsafeAccumArray' f initial (l,u) n@(I# n#) ies = runST (ST $ \s1# ->
     case newArray# n# initial s1#          of { (# s2#, marr# #) ->
-    foldr (adjust f marr#) (done l u n marr#) ies s2# })
+    foldr (adjust' f marr#) (done l u n marr#) ies s2# })
 
 {-# INLINE adjust #-}
 adjust :: (e -> a -> e) -> MutableArray# s e -> (Int, a) -> STRep s b -> STRep s b
@@ -678,6 +680,18 @@ adjust f marr# (I# i#, new) next
                     case writeArray# marr# i# (f old new) s2# of
                         s3# -> next s3#
 
+{-# INLINE adjust' #-}
+adjust' :: (e -> a -> e)
+        -> MutableArray# s e
+        -> (Int, a)
+        -> STRep s b -> STRep s b
+adjust' f marr# (I# i#, new) next
+  = \s1# -> case readArray# marr# i# s1# of
+                (# s2#, old #) ->
+                    let !combined = f old new
+                    in next (writeArray# marr# i# combined s2#)
+
+
 -- | Constructs an array identical to the first argument except that it has
 -- been updated by the associations in the right argument.
 -- For example, if @m@ is a 1-origin, @n@ by @n@ matrix, then
@@ -706,6 +720,8 @@ unsafeReplace arr ies = runST (do
 --
 -- > accumArray f z b = accum f (array b [(i, z) | i <- range b])
 --
+-- @accum@ is strict in all the results of applying the accumulation.
+-- However, it is lazy in the initial values of the array.
 {-# INLINE accum #-}
 accum :: Ix i => (e -> a -> e) -> Array i e -> [(i, a)] -> Array i e
 accum f arr@(Array l u n _) ies =
@@ -715,7 +731,7 @@ accum f arr@(Array l u n _) ies =
 unsafeAccum :: (e -> a -> e) -> Array i e -> [(Int, a)] -> Array i e
 unsafeAccum f arr ies = runST (do
     STArray l u n marr# <- thawSTArray arr
-    ST (foldr (adjust f marr#) (done l u n marr#) ies))
+    ST (foldr (adjust' f marr#) (done l u n marr#) ies))
 
 {-# INLINE [1] amap #-}  -- See Note [amap]
 amap :: (a -> b) -> Array i a -> Array i b