Add 'fromListWithKey' to HashMap (#246)

* Add 'fromListWithKey' to HashMap

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

* Define `unsafeInsertWith` using `unsafeInsertWithKey`

See #246 (comment)
for performance numbers.

* Improve documentation of fromListWithKey

Thanks to @emilypi for the suggestions

* Add a test for fromListWithKey

It models the test for fromListWith but makes
sure that values are combined in a way that depends
on the key.

* Clarify the documentation for fromListWithKey

* Improve properties for fromListWith and fromListWithKey

The old properties used associative operators to combine
values when there were duplicate keys.
With this diff we're using a non-commutative and
non-associative operator which can catch more bugs.

* Make sure to return an unboxed tuple in unsafeFromListWithKey

This got lost in a rebase

* Use the free magma to ensure that we preserve the order of operations

* Update fromListWithKey documentation to use non-commutative, non-associative operators

Since the combining function is applied in a way that can be
counter-intuitive it's more pedagical to to use operators which
better illustrate this behaviour.

Author: josefs

Data/HashMap/Base.hs

@@ -82,6 +82,7 @@ module Data.HashMap.Base
     , toList
     , fromList
     , fromListWith
+    , fromListWithKey
 
       -- Internals used by the strict version
     , Hash
@@ -1031,14 +1032,20 @@ insertModifyingArr x f k0 ary0 = go k0 ary0 0 (A.length ary0)
 unsafeInsertWith :: forall k v. (Eq k, Hashable k)
                  => (v -> v -> v) -> k -> v -> HashMap k v
                  -> HashMap k v
-unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
+unsafeInsertWith f k0 v0 m0 = unsafeInsertWithKey (const f) k0 v0 m0
+{-# INLINABLE unsafeInsertWith #-}
+
+unsafeInsertWithKey :: forall k v. (Eq k, Hashable k)
+                 => (k -> v -> v -> v) -> k -> v -> HashMap k v
+                 -> HashMap k v
+unsafeInsertWithKey f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
   where
     h0 = hash k0
     go :: Hash -> k -> v -> Shift -> HashMap k v -> ST s (HashMap k v)
     go !h !k x !_ Empty = return $! Leaf h (L k x)
     go h k x s t@(Leaf hy l@(L ky y))
         | hy == h = if ky == k
-                    then return $! Leaf h (L k (f x y))
+                    then return $! Leaf h (L k (f k x y))
                     else return $! collision h l (L k x)
         | otherwise = two s h k x hy t
     go h k x s t@(BitmapIndexed b ary)
@@ -1059,9 +1066,9 @@ unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
         return t
       where i = index h s
     go h k x s t@(Collision hy v)
-        | h == hy   = return $! Collision h (updateOrSnocWith (\a b -> (# f a b #)) k x v)
+        | h == hy   = return $! Collision h (updateOrSnocWithKey (\key a b -> (# f key a b #) ) k x v)
         | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
-{-# INLINABLE unsafeInsertWith #-}
+{-# INLINABLE unsafeInsertWithKey #-}
 
 -- | /O(log n)/ Remove the mapping for the specified key from this map
 -- if present.
@@ -1916,6 +1923,34 @@ fromListWith :: (Eq k, Hashable k) => (v -> v -> v) -> [(k, v)] -> HashMap k v
 fromListWith f = L.foldl' (\ m (k, v) -> unsafeInsertWith f k v m) empty
 {-# INLINE fromListWith #-}
 
+-- | /O(n*log n)/ Construct a map from a list of elements.  Uses
+-- the provided function to merge duplicate entries.
+--
+-- === Examples
+--
+-- Given a list of key-value pairs where the keys are of different flavours, e.g:
+--
+-- > data Key = Div | Sub
+--
+-- and the values need to be combined differently when there are duplicates,
+-- depending on the key:
+--
+-- > combine Div = div
+-- > combine Sub = (-)
+--
+-- then @fromListWithKey@ can be used as follows:
+--
+-- > fromListWithKey combine [(Div, 2), (Div, 6), (Sub, 2), (Sub, 3)]
+-- > = fromList [(Div, 3), (Sub, 1)]
+--
+-- More generally, duplicate entries are accumulated as follows;
+--
+-- > fromListWith f [(k, a), (k, b), (k, c), (k, d)]
+-- > = fromList [(k, f k d (f k c (f k b a)))]
+fromListWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> [(k, v)] -> HashMap k v
+fromListWithKey f = L.foldl' (\ m (k, v) -> unsafeInsertWithKey f k v m) empty
+{-# INLINE fromListWithKey #-}
+
 ------------------------------------------------------------------------
 -- Array operations
 

Data/HashMap/Lazy.hs

@@ -94,6 +94,7 @@ module Data.HashMap.Lazy
     , toList
     , fromList
     , fromListWith
+    , fromListWithKey
 
       -- ** HashSets
     , HS.keysSet

Data/HashMap/Strict.hs

@@ -93,6 +93,7 @@ module Data.HashMap.Strict
     , toList
     , fromList
     , fromListWith
+    , fromListWithKey
 
       -- ** HashSets
     , HS.keysSet

Data/HashMap/Strict/Base.hs

@@ -95,6 +95,7 @@ module Data.HashMap.Strict.Base
     , toList
     , fromList
     , fromListWith
+    , fromListWithKey
     ) where
 
 import Data.Bits ((.&.), (.|.))
@@ -109,7 +110,8 @@ import Prelude hiding (map, lookup)
 import qualified Data.HashMap.Array as A
 import qualified Data.HashMap.Base as HM
 import Data.HashMap.Base hiding (
-    alter, alterF, adjust, fromList, fromListWith, insert, insertWith,
+    alter, alterF, adjust, fromList, fromListWith, fromListWithKey,
+    insert, insertWith,
     differenceWith, intersectionWith, intersectionWithKey, map, mapWithKey,
     mapMaybe, mapMaybeWithKey, singleton, update, unionWith, unionWithKey,
     traverseWithKey)
@@ -189,13 +191,18 @@ insertWith f k0 v0 m0 = go h0 k0 v0 0 m0
 -- | In-place update version of insertWith
 unsafeInsertWith :: (Eq k, Hashable k) => (v -> v -> v) -> k -> v -> HashMap k v
                  -> HashMap k v
-unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
+unsafeInsertWith f k0 v0 m0 = unsafeInsertWithKey (const f) k0 v0 m0
+{-# INLINABLE unsafeInsertWith #-}
+
+unsafeInsertWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> k -> v -> HashMap k v
+                    -> HashMap k v
+unsafeInsertWithKey f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
   where
     h0 = hash k0
     go !h !k x !_ Empty = return $! leaf h k x
     go h k x s t@(Leaf hy l@(L ky y))
         | hy == h = if ky == k
-                    then return $! leaf h k (f x y)
+                    then return $! leaf h k (f k x y)
                     else do
                         let l' = x `seq` (L k x)
                         return $! collision h l l'
@@ -218,9 +225,9 @@ unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
         return t
       where i = index h s
     go h k x s t@(Collision hy v)
-        | h == hy   = return $! Collision h (updateOrSnocWith f k x v)
+        | h == hy   = return $! Collision h (updateOrSnocWithKey f k x v)
         | otherwise = go h k x s $ BitmapIndexed (mask hy s) (A.singleton t)
-{-# INLINABLE unsafeInsertWith #-}
+{-# INLINABLE unsafeInsertWithKey #-}
 
 -- | /O(log n)/ Adjust the value tied to a given key in this map only
 -- if it is present. Otherwise, leave the map alone.
@@ -639,6 +646,34 @@ fromListWith :: (Eq k, Hashable k) => (v -> v -> v) -> [(k, v)] -> HashMap k v
 fromListWith f = L.foldl' (\ m (k, v) -> unsafeInsertWith f k v m) empty
 {-# INLINE fromListWith #-}
 
+-- | /O(n*log n)/ Construct a map from a list of elements.  Uses
+-- the provided function to merge duplicate entries.
+--
+-- === Examples
+--
+-- Given a list of key-value pairs where the keys are of different flavours, e.g:
+--
+-- > data Key = Div | Sub
+--
+-- and the values need to be combined differently when there are duplicates,
+-- depending on the key:
+--
+-- > combine Div = div
+-- > combine Sub = (-)
+--
+-- then @fromListWithKey@ can be used as follows:
+--
+-- > fromListWithKey combine [(Div, 2), (Div, 6), (Sub, 2), (Sub, 3)]
+-- > = fromList [(Div, 3), (Sub, 1)]
+--
+-- More generally, duplicate entries are accumulated as follows;
+--
+-- > fromListWith f [(k, a), (k, b), (k, c), (k, d)]
+-- > = fromList [(k, f k d (f k c (f k b a)))]
+fromListWithKey :: (Eq k, Hashable k) => (k -> v -> v -> v) -> [(k, v)] -> HashMap k v
+fromListWithKey f = L.foldl' (\ m (k, v) -> unsafeInsertWithKey f k v m) empty
+{-# INLINE fromListWithKey #-}
+
 ------------------------------------------------------------------------
 -- Array operations
 

tests/HashMapProperties.hs

@@ -363,13 +363,31 @@ pFilterWithKey = M.filterWithKey p `eq_` HM.filterWithKey p
 ------------------------------------------------------------------------
 -- ** Conversions
 
+-- The free magma is used to test that operations are applied in the
+-- same order.
+data Magma a
+  = Leaf a
+  | Op (Magma a) (Magma a)
+  deriving (Show, Eq, Ord)
+
+instance Hashable a => Hashable (Magma a) where
+  hashWithSalt s (Leaf a) = hashWithSalt s (hashWithSalt (1::Int) a)
+  hashWithSalt s (Op m n) = hashWithSalt s (hashWithSalt (hashWithSalt (2::Int) m) n)
+
 -- 'eq_' already calls fromList.
 pFromList :: [(Key, Int)] -> Bool
 pFromList = id `eq_` id
 
 pFromListWith :: [(Key, Int)] -> Bool
-pFromListWith kvs = (M.toAscList $ M.fromListWith (+) kvs) ==
-                    (toAscList $ HM.fromListWith (+) kvs)
+pFromListWith kvs = (M.toAscList $ M.fromListWith Op kvsM) ==
+                    (toAscList $ HM.fromListWith Op kvsM)
+  where kvsM = fmap (fmap Leaf) kvs
+
+pFromListWithKey :: [(Key, Int)] -> Bool
+pFromListWithKey kvs = (M.toAscList $ M.fromListWithKey combine kvsM) ==
+                       (toAscList $ HM.fromListWithKey combine kvsM)
+  where kvsM = fmap (\(K k,v) -> (Leaf k, Leaf v)) kvs
+        combine k v1 v2 = Op k (Op v1 v2)
 
 pToList :: [(Key, Int)] -> Bool
 pToList = M.toAscList `eq` toAscList
@@ -467,6 +485,7 @@ tests =
       , testProperty "keys" pKeys
       , testProperty "fromList" pFromList
       , testProperty "fromListWith" pFromListWith
+      , testProperty "fromListWithKey" pFromListWithKey
       , testProperty "toList" pToList
       ]
     ]