@@ -107,6 +107,8 @@ module Data.HashMap.Base
107107 , insertModifying
108108 , ptrEq
109109 , adjust #
110+ , unionWithKey #
111+ , unsafeInsertWith
110112 ) where
111113
112114#if __GLASGOW_HASKELL__ < 710
@@ -650,7 +652,7 @@ insert' h0 k0 v0 m0 = go h0 k0 v0 0 m0
650652 else Full (update16 ary i st')
651653 where i = index h s
652654 go h k x s t@ (Collision hy v)
653- | h == hy = Collision h (updateOrSnocWith const k x v)
655+ | h == hy = Collision h (updateOrSnocWith ( \ v1 _ -> ( # v1 # )) k x v)
654656 | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A. singleton t)
655657{-# INLINABLE insert' #-}
656658
@@ -773,7 +775,7 @@ unsafeInsert k0 v0 m0 = runST (go h0 k0 v0 0 m0)
773775 return t
774776 where i = index h s
775777 go h k x s t@ (Collision hy v)
776- | h == hy = return $! Collision h (updateOrSnocWith const k x v)
778+ | h == hy = return $! Collision h (updateOrSnocWith ( \ v1 _ -> ( # v1 # )) k x v)
777779 | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A. singleton t)
778780{-# INLINABLE unsafeInsert #-}
779781
@@ -882,7 +884,7 @@ insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0)
882884
883885-- | In-place update version of insertWith
884886unsafeInsertWith :: forall k v . (Eq k , Hashable k )
885- => (v -> v -> v ) -> k -> v -> HashMap k v
887+ => (v -> v -> ( # v # ) ) -> k -> v -> HashMap k v
886888 -> HashMap k v
887889unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
888890 where
@@ -891,7 +893,7 @@ unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
891893 go ! h ! k x ! _ Empty = return $! Leaf h (L k x)
892894 go h k x s (Leaf hy l@ (L ky y))
893895 | hy == h = if ky == k
894- then return $! Leaf h (L k (f x y) )
896+ then case f x y of ( # v # ) -> return $! Leaf h (L k v )
895897 else return $! collision h l (L k x)
896898 | otherwise = two s h k x hy ky y
897899 go h k x s t@ (BitmapIndexed b ary)
@@ -1256,22 +1258,27 @@ unionWith f = unionWithKey (const f)
12561258-- result.
12571259unionWithKey :: (Eq k , Hashable k ) => (k -> v -> v -> v ) -> HashMap k v -> HashMap k v
12581260 -> HashMap k v
1259- unionWithKey f = go 0
1261+ unionWithKey f m = unionWithKey# (\ k v1 v2 -> (# f k v1 v2 # )) m
1262+ {-# INLINE unionWithKey #-}
1263+
1264+ unionWithKey# :: (Eq k , Hashable k ) => (k -> v -> v -> (# v # )) -> HashMap k v -> HashMap k v
1265+ -> HashMap k v
1266+ unionWithKey# f = go 0
12601267 where
12611268 -- empty vs. anything
12621269 go ! _ t1 Empty = t1
12631270 go _ Empty t2 = t2
12641271 -- leaf vs. leaf
12651272 go s t1@ (Leaf h1 l1@ (L k1 v1)) t2@ (Leaf h2 l2@ (L k2 v2))
12661273 | h1 == h2 = if k1 == k2
1267- then Leaf h1 ( L k1 ( f k1 v1 v2) )
1274+ then case f k1 v1 v2 of ( # v # ) -> Leaf h1 ( L k1 v )
12681275 else collision h1 l1 l2
12691276 | otherwise = goDifferentHash s h1 h2 t1 t2
12701277 go s t1@ (Leaf h1 (L k1 v1)) t2@ (Collision h2 ls2)
12711278 | h1 == h2 = Collision h1 (updateOrSnocWithKey f k1 v1 ls2)
12721279 | otherwise = goDifferentHash s h1 h2 t1 t2
12731280 go s t1@ (Collision h1 ls1) t2@ (Leaf h2 (L k2 v2))
1274- | h1 == h2 = Collision h1 (updateOrSnocWithKey (flip . f ) k2 v2 ls1)
1281+ | h1 == h2 = Collision h1 (updateOrSnocWithKey (\ q w x -> f q x w ) k2 v2 ls1)
12751282 | otherwise = goDifferentHash s h1 h2 t1 t2
12761283 go s t1@ (Collision h1 ls1) t2@ (Collision h2 ls2)
12771284 | h1 == h2 = Collision h1 (updateOrConcatWithKey f ls1 ls2)
@@ -1336,7 +1343,8 @@ unionWithKey f = go 0
13361343 where
13371344 m1 = mask h1 s
13381345 m2 = mask h2 s
1339- {-# INLINE unionWithKey #-}
1346+ {-# INLINE unionWithKey# #-}
1347+
13401348
13411349-- | Strict in the result of @f@.
13421350unionArrayBy :: (a -> a -> a ) -> Bitmap -> Bitmap -> A. Array a -> A. Array a
@@ -1667,7 +1675,7 @@ fromList = L.foldl' (\ m (k, v) -> unsafeInsert k v m) empty
16671675-- | /O(n*log n)/ Construct a map from a list of elements. Uses
16681676-- the provided function to merge duplicate entries.
16691677fromListWith :: (Eq k , Hashable k ) => (v -> v -> v ) -> [(k , v )] -> HashMap k v
1670- fromListWith f = L. foldl' (\ m (k, v) -> unsafeInsertWith f k v m) empty
1678+ fromListWith f = L. foldl' (\ m (k, v) -> unsafeInsertWith ( \ x y -> ( # f x y # )) k v m) empty
16711679{-# INLINE fromListWith #-}
16721680
16731681------------------------------------------------------------------------
@@ -1719,12 +1727,12 @@ updateWith# f k0 ary0 = go k0 ary0 0 (A.length ary0)
17191727 | otherwise -> go k ary (i+ 1 ) n
17201728{-# INLINABLE updateWith# #-}
17211729
1722- updateOrSnocWith :: Eq k => (v -> v -> v ) -> k -> v -> A. Array (Leaf k v )
1730+ updateOrSnocWith :: Eq k => (v -> v -> ( # v # ) ) -> k -> v -> A. Array (Leaf k v )
17231731 -> A. Array (Leaf k v )
17241732updateOrSnocWith f = updateOrSnocWithKey (const f)
17251733{-# INLINABLE updateOrSnocWith #-}
17261734
1727- updateOrSnocWithKey :: Eq k => (k -> v -> v -> v ) -> k -> v -> A. Array (Leaf k v )
1735+ updateOrSnocWithKey :: Eq k => (k -> v -> v -> ( # v # ) ) -> k -> v -> A. Array (Leaf k v )
17281736 -> A. Array (Leaf k v )
17291737updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A. length ary0)
17301738 where
@@ -1736,15 +1744,15 @@ updateOrSnocWithKey f k0 v0 ary0 = go k0 v0 ary0 0 (A.length ary0)
17361744 A. write mary n (L k v)
17371745 return mary
17381746 | otherwise = case A. index ary i of
1739- (L kx y) | k == kx -> A. update ary i (L k (f k v y) )
1747+ (L kx y) | k == kx -> case f k v y of ( # y' # ) -> A. update ary i (L k y' )
17401748 | otherwise -> go k v ary (i+ 1 ) n
17411749{-# INLINABLE updateOrSnocWithKey #-}
17421750
1743- updateOrConcatWith :: Eq k => (v -> v -> v ) -> A. Array (Leaf k v ) -> A. Array (Leaf k v ) -> A. Array (Leaf k v )
1751+ updateOrConcatWith :: Eq k => (v -> v -> ( # v # ) ) -> A. Array (Leaf k v ) -> A. Array (Leaf k v ) -> A. Array (Leaf k v )
17441752updateOrConcatWith f = updateOrConcatWithKey (const f)
17451753{-# INLINABLE updateOrConcatWith #-}
17461754
1747- updateOrConcatWithKey :: Eq k => (k -> v -> v -> v ) -> A. Array (Leaf k v ) -> A. Array (Leaf k v ) -> A. Array (Leaf k v )
1755+ updateOrConcatWithKey :: Eq k => (k -> v -> v -> ( # v # ) ) -> A. Array (Leaf k v ) -> A. Array (Leaf k v ) -> A. Array (Leaf k v )
17481756updateOrConcatWithKey f ary1 ary2 = A. run $ do
17491757 -- first: look up the position of each element of ary2 in ary1
17501758 let indices = A. map (\ (L k _) -> indexOf k ary1) ary2
@@ -1763,7 +1771,7 @@ updateOrConcatWithKey f ary1 ary2 = A.run $ do
17631771 Just i1 -> do -- key occurs in both arrays, store combination in position i1
17641772 L k v1 <- A. indexM ary1 i1
17651773 L _ v2 <- A. indexM ary2 i2
1766- A. write mary i1 (L k (f k v1 v2) )
1774+ case f k v1 v2 of ( # v' # ) -> A. write mary i1 (L k v' )
17671775 go iEnd (i2+ 1 )
17681776 Nothing -> do -- key is only in ary2, append to end
17691777 A. write mary iEnd =<< A. indexM ary2 i2
0 commit comments