Traverse arrays better

The previous approach was to convert the array to a list, traverse
the list, and then put it back into an array. I figured it might
be worth trying a more direct approach. It seems to work
rather better. I also added rules to use a very direct implementation
for `ST` and `IO`.

Author: treeowl

Data/HashMap/Array.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples #-}
+{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples, ScopedTypeVariables #-}
 {-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-}
 
 -- | Zero based arrays.
@@ -57,8 +57,9 @@ import Control.Applicative (Applicative (..), (<$>))
 #endif
 import Control.Applicative (liftA2)
 import Control.DeepSeq
-import GHC.Exts(Int(..), Int#, reallyUnsafePtrEquality#, tagToEnum#, unsafeCoerce#, State#)
+import GHC.Exts(Int(..), Int#, reallyUnsafePtrEquality#, tagToEnum#, unsafeCoerce#, State#, (+#))
 import GHC.ST (ST(..))
+import Control.Monad.ST (stToIO)
 
 #if __GLASGOW_HASKELL__ >= 709
 import Prelude hiding (filter, foldr, length, map, read, traverse)
@@ -475,49 +476,36 @@ fromList n xs0 =
 toList :: Array a -> [a]
 toList = foldr (:) []
 
-data SList a = SCons !a (SList a) | SNil
+newtype STA a = STA {_runSTA :: forall s. MutableArray# s a -> ST s (Array a)}
 
-traverseToSList
-  :: Applicative f
-  => (a -> f b) -> [a] -> f (SList b)
-traverseToSList f = go
-  where
-    go (a : as) = liftA2 SCons (f a) (go as)
-    go [] = pure SNil
-
-_slength :: SList a -> Int
-_slength = go 0 where
-  go !acc SNil = acc
-  go acc (SCons _ xs) = go (acc + 1) xs
+runSTA :: Int -> STA a -> Array a
+runSTA !n (STA m) = runST $ new_ n >>= \ (MArray ar) -> m ar
 
-fromSList :: Int -> SList a -> Array a
-fromSList n xs0 =
-    CHECK_EQ("fromSList", n, _slength xs0)
-        run $ do
-            mary <- new_ n
-            go xs0 mary 0
+traverse :: forall f a b. Applicative f => (a -> f b) -> Array a -> f (Array b)
+traverse f = \ !ary -> runSTA (length ary) <$> foldr go stop ary 0#
   where
-    go SNil !mary !_ = return mary
-    go (SCons x xs) mary i = do write mary i x
-                                go xs mary (i + 1)
-
-traverse :: Applicative f => (a -> f b) -> Array a -> f (Array b)
-traverse f = \ ary -> fromList (length ary) `fmap`
-                      Traversable.traverse f (toList ary)
+    go :: a -> (Int# -> f (STA b)) -> Int# -> f (STA b)
+    go a r i = liftA2 (\b (STA m) -> STA $ \mry# -> write (MArray mry#) (I# i) b >> m mry#) (f a) (r (i +# 1#))
+    stop :: Int# -> f (STA b)
+    stop _i = pure (STA (\mry# -> unsafeFreeze (MArray mry#)))
 {-# INLINE [1] traverse #-}
 
-traverse' :: Applicative f => (a -> f b) -> Array a -> f (Array b)
-traverse' f = \ary -> fromSList (length ary) `fmap`
-                      traverseToSList f (toList ary)
+traverse' :: forall f a b. Applicative f => (a -> f b) -> Array a -> f (Array b)
+traverse' f = \ !ary -> runSTA (length ary) <$> foldr go stop ary 0#
+  where
+    go :: a -> (Int# -> f (STA b)) -> Int# -> f (STA b)
+    go a r i = liftA2 (\ !b (STA m) -> STA $ \mry# -> write (MArray mry#) (I# i) b >> m mry#) (f a) (r (i +# 1#))
+    stop :: Int# -> f (STA b)
+    stop _i = pure (STA (\mry# -> unsafeFreeze (MArray mry#)))
 {-# INLINE [1] traverse' #-}
 
--- Traversing in ST, we don't need to make a list; we
+-- Traversing in ST, we don't need to get fancy; we
 -- can just do it directly.
 traverseST :: (a -> ST s b) -> Array a -> ST s (Array b)
 traverseST f = \ ary0 ->
   let
     !len = length ary0
-    go k mary
+    go k !mary
       | k == len = return mary
       | otherwise = do
           x <- indexM ary0 k
@@ -527,17 +515,33 @@ traverseST f = \ ary0 ->
   in new_ len >>= (go 0 >=> unsafeFreeze)
 {-# INLINE traverseST #-}
 
+traverseIO :: (a -> IO b) -> Array a -> IO (Array b)
+traverseIO f = \ ary0 ->
+  let
+    !len = length ary0
+    go k !mary
+      | k == len = return mary
+      | otherwise = do
+          x <- stToIO $ indexM ary0 k
+          y <- f x
+          stToIO $ write mary k y
+          go (k + 1) mary
+  in stToIO (new_ len) >>= (go 0 >=> stToIO . unsafeFreeze)
+{-# INLINE traverseIO #-}
+
+
 {-# RULES
 "traverse/ST" forall f. traverse f = traverseST f
+"traverse/IO" forall f. traverse f = traverseIO f
  #-}
 
--- Traversing in ST, we don't need to make a list; we
+-- Traversing in ST, we don't need to get fancy; we
 -- can just do it directly.
 traverseST' :: (a -> ST s b) -> Array a -> ST s (Array b)
 traverseST' f = \ ary0 ->
   let
     !len = length ary0
-    go k mary
+    go k !mary
       | k == len = return mary
       | otherwise = do
           x <- indexM ary0 k
@@ -547,8 +551,23 @@ traverseST' f = \ ary0 ->
   in new_ len >>= (go 0 >=> unsafeFreeze)
 {-# INLINE traverseST' #-}
 
+traverseIO' :: (a -> IO b) -> Array a -> IO (Array b)
+traverseIO' f = \ ary0 ->
+  let
+    !len = length ary0
+    go k !mary
+      | k == len = return mary
+      | otherwise = do
+          x <- stToIO $ indexM ary0 k
+          !y <- f x
+          stToIO $ write mary k y
+          go (k + 1) mary
+  in stToIO (new_ len) >>= (go 0 >=> stToIO . unsafeFreeze)
+{-# INLINE traverseIO' #-}
+
 {-# RULES
 "traverse'/ST" forall f. traverse' f = traverseST' f
+"traverse'/IO" forall f. traverse' f = traverseIO' f
  #-}
 
 filter :: (a -> Bool) -> Array a -> Array a

Data/HashMap/Strict/Base.hs

@@ -93,7 +93,7 @@ module Data.HashMap.Strict.Base
 import Data.Bits ((.&.), (.|.))
 
 #if !MIN_VERSION_base(4,8,0)
-import Data.Functor((<$>))
+import Control.Applicative (Applicative (..), (<$>))
 #endif
 import qualified Data.List as L
 import Data.Hashable (Hashable)