Make Strict.traverseWithKey actually strict

Fixes #192

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

@@ -46,14 +46,16 @@ 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.ST (ST(..))
@@ -473,11 +475,42 @@ fromList n xs0 =
 toList :: Array a -> [a]
 toList = foldr (:) []
 
+data SList a = SCons !a (SList a) | SNil
+
+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
+
+fromSList :: Int -> SList a -> Array a
+fromSList n xs0 =
+    CHECK_EQ("fromSList", n, _slength xs0)
+        run $ do
+            mary <- new_ n
+            go xs0 mary 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)
 {-# 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)
+{-# INLINE [1] traverse' #-}
+
 -- Traversing in ST, we don't need to make a list; we
 -- can just do it directly.
 traverseST :: (a -> ST s b) -> Array a -> ST s (Array b)
@@ -498,6 +531,26 @@ traverseST f = \ ary0 ->
 "traverse/ST" forall f. traverse f = traverseST f
  #-}
 
+-- Traversing in ST, we don't need to make a list; 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' #-}
+
+{-# RULES
+"traverse'/ST" forall f. traverse' f = traverseST' f
+ #-}
+
 filter :: (a -> Bool) -> Array a -> Array a
 filter p = \ ary ->
     let !n = length ary

Data/HashMap/Base.hs

@@ -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

@@ -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