Merge pull request #194 from treeowl/strict-traverse

Make Strict.traverseWithKey actually strict

Author: treeowl

CHANGES.md

@@ -4,6 +4,9 @@
 
  * Add `HashMap.keysSet`.
 
+ * Make `HashMap.Strict.traverseWithKey` force the results before
+   installing them in the map.
+
 ## 0.2.9.0
 
  * Add `Ord/Ord1/Ord2` instances. (Thanks, Oleg Grenrus)

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.
@@ -46,17 +46,20 @@ module Data.HashMap.Array
     , map
     , map'
     , traverse
+    , traverse'
     , filter
     , toList
+    , fromList
     ) where
 
-import qualified Data.Traversable as Traversable
-#if __GLASGOW_HASKELL__ < 709
-import Control.Applicative (Applicative)
+#if !MIN_VERSION_base(4,8,0)
+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)
@@ -473,18 +476,36 @@ fromList n xs0 =
 toList :: Array a -> [a]
 toList = foldr (:) []
 
-traverse :: Applicative f => (a -> f b) -> Array a -> f (Array b)
-traverse f = \ ary -> fromList (length ary) `fmap`
-                      Traversable.traverse f (toList ary)
+newtype STA a = STA {_runSTA :: forall s. MutableArray# s a -> ST s (Array a)}
+
+runSTA :: Int -> STA a -> Array a
+runSTA !n (STA m) = runST $ new_ n >>= \ (MArray ar) -> m ar
+
+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
+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 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
@@ -494,8 +515,59 @@ 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 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
+      | k == len = return mary
+      | otherwise = do
+          x <- indexM ary0 k
+          !y <- f x
+          write mary k y
+          go (k + 1) mary
+  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/Base.hs

@@ -634,7 +634,7 @@ infixl 9 !
 
 -- | Create a 'Collision' value with two 'Leaf' values.
 collision :: Hash -> Leaf k v -> Leaf k v -> HashMap k v
-collision h e1 e2 =
+collision h !e1 !e2 =
     let v = A.run $ do mary <- A.new 2 e1
                        A.write mary 1 e2
                        return mary
@@ -1432,18 +1432,25 @@ map f = mapWithKey (const f)
 -- TODO: We should be able to use mutation to create the new
 -- 'HashMap'.
 
--- | /O(n)/ Transform this map by accumulating an Applicative result
--- from every value.
-traverseWithKey :: Applicative f => (k -> v1 -> f v2) -> HashMap k v1
-                -> f (HashMap k v2)
+-- | /O(n)/ Perform an 'Applicative' action for each key-value pair
+-- in a 'HashMap' and produce a 'HashMap' of all the results.
+--
+-- Note: the order in which the actions occur is unspecified. In particular,
+-- when the map contains hash collisions, the order in which the actions
+-- associated with the keys involved will depend in an unspecified way on
+-- their insertion order.
+traverseWithKey
+  :: Applicative f
+  => (k -> v1 -> f v2)
+  -> HashMap k v1 -> f (HashMap k v2)
 traverseWithKey f = go
   where
     go Empty                 = pure Empty
     go (Leaf h (L k v))      = Leaf h . L k <$> f k v
     go (BitmapIndexed b ary) = BitmapIndexed b <$> A.traverse go ary
     go (Full ary)            = Full <$> A.traverse go ary
     go (Collision h ary)     =
-        Collision h <$> A.traverse (\ (L k v) -> L k <$> f k v) ary
+        Collision h <$> A.traverse' (\ (L k v) -> L k <$> f k v) ary
 {-# INLINE traverseWithKey #-}
 
 ------------------------------------------------------------------------

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)
@@ -104,7 +104,8 @@ import qualified Data.HashMap.Base as HM
 import Data.HashMap.Base hiding (
     alter, alterF, adjust, fromList, fromListWith, insert, insertWith,
     differenceWith, intersectionWith, intersectionWithKey, map, mapWithKey,
-    mapMaybe, mapMaybeWithKey, singleton, update, unionWith, unionWithKey)
+    mapMaybe, mapMaybeWithKey, singleton, update, unionWith, unionWithKey,
+    traverseWithKey)
 import Data.HashMap.Unsafe (runST)
 #if MIN_VERSION_base(4,8,0)
 import Data.Functor.Identity
@@ -522,8 +523,31 @@ mapMaybe :: (v1 -> Maybe v2) -> HashMap k v1 -> HashMap k v2
 mapMaybe f = mapMaybeWithKey (const f)
 {-# INLINE mapMaybe #-}
 
-
--- TODO: Should we add a strict traverseWithKey?
+-- | /O(n)/ Perform an 'Applicative' action for each key-value pair
+-- in a 'HashMap' and produce a 'HashMap' of all the results. Each 'HashMap'
+-- will be strict in all its values.
+--
+-- @
+-- traverseWithKey f = fmap ('map' id) . "Data.HashMap.Lazy".'Data.HashMap.Lazy.traverseWithKey' f
+-- @
+--
+-- Note: the order in which the actions occur is unspecified. In particular,
+-- when the map contains hash collisions, the order in which the actions
+-- associated with the keys involved will depend in an unspecified way on
+-- their insertion order.
+traverseWithKey
+  :: Applicative f
+  => (k -> v1 -> f v2)
+  -> HashMap k v1 -> f (HashMap k v2)
+traverseWithKey f = go
+  where
+    go Empty                 = pure Empty
+    go (Leaf h (L k v))      = leaf h k <$> f k v
+    go (BitmapIndexed b ary) = BitmapIndexed b <$> A.traverse' go ary
+    go (Full ary)            = Full <$> A.traverse' go ary
+    go (Collision h ary)     =
+        Collision h <$> A.traverse' (\ (L k v) -> (L k $!) <$> f k v) ary
+{-# INLINE traverseWithKey #-}
 
 ------------------------------------------------------------------------
 -- * Difference and intersection
@@ -643,5 +667,5 @@ updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
 -- inserted into the constructor.
 
 leaf :: Hash -> k -> v -> HashMap k v
-leaf h k !v = Leaf h (L k v)
+leaf h k = \ !v -> Leaf h (L k v)
 {-# INLINE leaf #-}

tests/HashMapProperties.hs

@@ -246,6 +246,11 @@ pUnions xss = M.toAscList (M.unions (map M.fromList xss)) ==
 pMap :: [(Key, Int)] -> Bool
 pMap = M.map (+ 1) `eq_` HM.map (+ 1)
 
+pTraverse :: [(Key, Int)] -> Bool
+pTraverse xs =
+  L.sort (fmap (L.sort . M.toList) (M.traverseWithKey (\_ v -> [v + 1, v + 2]) (M.fromList (take 10 xs))))
+     == L.sort (fmap (L.sort . HM.toList) (HM.traverseWithKey (\_ v -> [v + 1, v + 2]) (HM.fromList (take 10 xs))))
+
 ------------------------------------------------------------------------
 -- ** Difference and intersection
 
@@ -382,6 +387,7 @@ tests =
     , testProperty "unions" pUnions
     -- Transformations
     , testProperty "map" pMap
+    , testProperty "traverse" pTraverse
     -- Folds
     , testGroup "folds"
       [ testProperty "foldr" pFoldr