Add HashMap.adjustWithKey.

CHANGES.md

@@ -4,6 +4,7 @@
  * Remove support for GHC versions before 7.8. (Thanks, Dmitry Ivanov!)
  * Use `SmallArray#` instead of `Array#` for GHC versions 7.10 and above.
    (Thanks, Dmitry Ivanov!)
+ * Add `HashMap.adjustWithKey`.
 
 ## 0.2.8.0
 

Data/HashMap/Base.hs

@@ -26,6 +26,7 @@ module Data.HashMap.Base
     , unsafeInsert
     , delete
     , adjust
+    , adjustWithKey
     , update
     , alter
 
@@ -713,12 +714,18 @@ delete k0 m0 = go h0 k0 0 m0
 -- | /O(log n)/ Adjust the value tied to a given key in this map only
 -- if it is present. Otherwise, leave the map alone.
 adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v
-adjust f k0 m0 = go h0 k0 0 m0
+adjust f = adjustWithKey (const f)
+{-# INLINABLE adjust #-}
+
+-- | /O(log n)/ Adjust the value tied to a given key in this map only
+-- if it is present. Otherwise, leave the map alone.
+adjustWithKey :: (Eq k, Hashable k) => (k -> v -> v) -> k -> HashMap k v -> HashMap k v
+adjustWithKey f k0 m0 = go h0 k0 0 m0
   where
     h0 = hash k0
     go !_ !_ !_ Empty = Empty
     go h k _ t@(Leaf hy (L ky y))
-        | hy == h && ky == k = Leaf h (L k (f y))
+        | hy == h && ky == k = Leaf h (L k (f k y))
         | otherwise          = t
     go h k s t@(BitmapIndexed b ary)
         | b .&. m == 0 = t
@@ -735,9 +742,9 @@ adjust f k0 m0 = go h0 k0 0 m0
             ary' = update16 ary i $! st'
         in Full ary'
     go h k _ t@(Collision hy v)
-        | h == hy   = Collision h (updateWith f k v)
+        | h == hy   = Collision h (updateWithKey f k v)
         | otherwise = t
-{-# INLINABLE adjust #-}
+{-# INLINABLE adjustWithKey #-}
 
 -- | /O(log n)/  The expression (@'update' f k map@) updates the value @x@ at @k@,
 -- (if it is in the map). If (f k x) is @'Nothing', the element is deleted.
@@ -1223,14 +1230,18 @@ indexOf k0 ary0 = go k0 ary0 0 (A.length ary0)
 {-# INLINABLE indexOf #-}
 
 updateWith :: Eq k => (v -> v) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v)
-updateWith f k0 ary0 = go k0 ary0 0 (A.length ary0)
+updateWith f = updateWithKey (const f)
+{-# INLINABLE updateWith #-}
+
+updateWithKey :: Eq k => (k -> v -> v) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v)
+updateWithKey f k0 ary0 = go k0 ary0 0 (A.length ary0)
   where
     go !k !ary !i !n
         | i >= n    = ary
         | otherwise = case A.index ary i of
-            (L kx y) | k == kx   -> A.update ary i (L k (f y))
+            (L kx y) | k == kx   -> A.update ary i (L k (f k y))
                      | otherwise -> go k ary (i+1) n
-{-# INLINABLE updateWith #-}
+{-# INLINABLE updateWithKey #-}
 
 updateOrSnocWith :: Eq k => (v -> v -> v) -> k -> v -> A.Array (Leaf k v)
                  -> A.Array (Leaf k v)

Data/HashMap/Lazy.hs

@@ -44,6 +44,7 @@ module Data.HashMap.Lazy
     , insertWith
     , delete
     , adjust
+    , adjustWithKey
     , update
     , alter
 

Data/HashMap/Strict.hs

@@ -44,6 +44,7 @@ module Data.HashMap.Strict
     , insertWith
     , delete
     , adjust
+    , adjustWithKey
     , update
     , alter
 
@@ -96,9 +97,9 @@ import Prelude hiding (map)
 import qualified Data.HashMap.Array as A
 import qualified Data.HashMap.Base as HM
 import Data.HashMap.Base hiding (
-    alter, adjust, fromList, fromListWith, insert, insertWith, differenceWith,
-    intersectionWith, intersectionWithKey, map, mapWithKey, mapMaybe,
-    mapMaybeWithKey, singleton, update, unionWith, unionWithKey)
+    alter, adjust, adjustWithKey, fromList, fromListWith, insert, insertWith,
+    differenceWith, intersectionWith, intersectionWithKey, map, mapWithKey,
+    mapMaybe, mapMaybeWithKey, singleton, update, unionWith, unionWithKey)
 import Data.HashMap.Unsafe (runST)
 
 -- $strictness
@@ -206,12 +207,18 @@ unsafeInsertWith f k0 v0 m0 = runST (go h0 k0 v0 0 m0)
 -- | /O(log n)/ Adjust the value tied to a given key in this map only
 -- if it is present. Otherwise, leave the map alone.
 adjust :: (Eq k, Hashable k) => (v -> v) -> k -> HashMap k v -> HashMap k v
-adjust f k0 m0 = go h0 k0 0 m0
+adjust f = adjustWithKey (const f)
+{-# INLINABLE adjust #-}
+
+-- | /O(log n)/ Adjust the value tied to a given key in this map only
+-- if it is present. Otherwise, leave the map alone.
+adjustWithKey :: (Eq k, Hashable k) => (k -> v -> v) -> k -> HashMap k v -> HashMap k v
+adjustWithKey f k0 m0 = go h0 k0 0 m0
   where
     h0 = hash k0
     go !_ !_ !_ Empty = Empty
     go h k _ t@(Leaf hy (L ky y))
-        | hy == h && ky == k = leaf h k (f y)
+        | hy == h && ky == k = leaf h k (f k y)
         | otherwise          = t
     go h k s t@(BitmapIndexed b ary)
         | b .&. m == 0 = t
@@ -228,9 +235,9 @@ adjust f k0 m0 = go h0 k0 0 m0
             ary' = update16 ary i $! st'
         in Full ary'
     go h k _ t@(Collision hy v)
-        | h == hy   = Collision h (updateWith f k v)
+        | h == hy   = Collision h (updateWithKey f k v)
         | otherwise = t
-{-# INLINABLE adjust #-}
+{-# INLINABLE adjustWithKey #-}
 
 -- | /O(log n)/  The expression (@'update' f k map@) updates the value @x@ at @k@,
 -- (if it is in the map). If (f k x) is @'Nothing', the element is deleted.
@@ -459,14 +466,18 @@ fromListWith f = L.foldl' (\ m (k, v) -> unsafeInsertWith f k v m) empty
 -- Array operations
 
 updateWith :: Eq k => (v -> v) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v)
-updateWith f k0 ary0 = go k0 ary0 0 (A.length ary0)
+updateWith f = updateWithKey (const f)
+{-# INLINABLE updateWith #-}
+
+updateWithKey :: Eq k => (k -> v -> v) -> k -> A.Array (Leaf k v) -> A.Array (Leaf k v)
+updateWithKey f k0 ary0 = go k0 ary0 0 (A.length ary0)
   where
     go !k !ary !i !n
         | i >= n    = ary
         | otherwise = case A.index ary i of
-            (L kx y) | k == kx   -> let !v' = f y in A.update ary i (L k v')
+            (L kx y) | k == kx   -> let !v' = f k y in A.update ary i (L k v')
                      | otherwise -> go k ary (i+1) n
-{-# INLINABLE updateWith #-}
+{-# INLINABLE updateWithKey #-}
 
 -- | Append the given key and value to the array. If the key is
 -- already present, instead update the value of the key by applying

tests/HashMapProperties.hs

@@ -17,17 +17,26 @@ import qualified Data.HashMap.Strict as HM
 import qualified Data.HashMap.Lazy as HM
 #endif
 import qualified Data.Map as M
-import Test.QuickCheck (Arbitrary, Property, (==>), (===))
+import Test.QuickCheck (Arbitrary, CoArbitrary, Property, (==>), (===))
+import Test.QuickCheck.Function (Fun(Fun), Function(function), functionMap)
 import Test.Framework (Test, defaultMain, testGroup)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
 
 -- Key type that generates more hash collisions.
 newtype Key = K { unK :: Int }
-            deriving (Arbitrary, Eq, Ord, Read, Show)
+            deriving (Arbitrary, CoArbitrary, Eq, Ord, Read, Show)
 
 instance Hashable Key where
     hashWithSalt salt k = hashWithSalt salt (unK k) `mod` 20
 
+instance Function Key where
+    function = functionMap unK K
+
+-- | Extracts the value of a ternary function.
+-- Copied from Test.QuickCheck.Function.applyFun3
+applyFun2 :: Fun (a, b) c -> (a -> b -> c)
+applyFun2 (Fun _ f) a b = f (a, b)
+
 ------------------------------------------------------------------------
 -- * Properties
 
@@ -156,6 +165,10 @@ pAdjust k = M.adjust succ k `eq_` HM.adjust succ k
 pUpdateAdjust :: Key -> [(Key, Int)] -> Bool
 pUpdateAdjust k = M.update (Just . succ) k `eq_` HM.update (Just . succ) k
 
+pAdjustWithKey :: Fun (Key, Int) Int -> Key -> [(Key, Int)] -> Bool
+pAdjustWithKey f k =
+    M.adjustWithKey (applyFun2 f) k `eq_` HM.adjustWithKey (applyFun2 f) k
+
 pUpdateDelete :: Key -> [(Key, Int)] -> Bool
 pUpdateDelete k = M.update (const Nothing) k `eq_` HM.update (const Nothing) k
 
@@ -313,6 +326,7 @@ tests =
       , testProperty "insertWith" pInsertWith
       , testProperty "adjust" pAdjust
       , testProperty "updateAdjust" pUpdateAdjust
+      , testProperty "adjustWithKey" pAdjustWithKey
       , testProperty "updateDelete" pUpdateDelete
       , testProperty "alterAdjust" pAlterAdjust
       , testProperty "alterInsert" pAlterInsert

tests/Strictness.hs

@@ -7,8 +7,8 @@ import Data.Hashable (Hashable(hashWithSalt))
 import Test.ChasingBottoms.IsBottom
 import Test.Framework (Test, defaultMain, testGroup)
 import Test.Framework.Providers.QuickCheck2 (testProperty)
-import Test.QuickCheck (Arbitrary(arbitrary), Property, (===), (.&&.))
-import Test.QuickCheck.Function
+import Test.QuickCheck (Arbitrary(arbitrary), CoArbitrary, Property, (===), (.&&.))
+import Test.QuickCheck.Function (Fun(Fun), Function(function), functionMap)
 import Test.QuickCheck.Poly (A)
 import Data.Maybe (fromMaybe, isJust)
 import Control.Arrow (second)
@@ -25,11 +25,14 @@ import qualified Data.HashMap.Strict as HM
 
 -- Key type that generates more hash collisions.
 newtype Key = K { unK :: Int }
-            deriving (Arbitrary, Eq, Ord, Show)
+            deriving (Arbitrary, CoArbitrary, Eq, Ord, Show)
 
 instance Hashable Key where
     hashWithSalt salt k = hashWithSalt salt (unK k) `mod` 20
 
+instance Function Key where
+    function = functionMap unK K
+
 instance (Arbitrary k, Arbitrary v, Eq k, Hashable k) =>
          Arbitrary (HashMap k v) where
     arbitrary = HM.fromList `fmap` arbitrary
@@ -40,6 +43,11 @@ instance Show (Int -> Int) where
 instance Show (Int -> Int -> Int) where
     show _ = "<function>"
 
+-- | Extracts the value of a ternary function.
+-- Copied from Test.QuickCheck.Function.applyFun3
+applyFun2 :: Fun (a, b) c -> (a -> b -> c)
+applyFun2 (Fun _ f) a b = f (a, b)
+
 ------------------------------------------------------------------------
 -- * Properties
 
@@ -65,6 +73,16 @@ pAdjustValueStrict k m
         []     -> True
         (k':_) -> isBottom $ HM.adjust (const bottom) k' m
 
+pAdjustWithKeyKeyStrict :: Fun (Key, Int) Int -> HashMap Key Int -> Bool
+pAdjustWithKeyKeyStrict f m = isBottom $ HM.adjustWithKey (applyFun2 f) bottom m
+
+pAdjustWithKeyValueStrict :: Key -> HashMap Key Int -> Bool
+pAdjustWithKeyValueStrict k m
+    | k `HM.member` m = isBottom $ HM.adjust (const bottom) k m
+    | otherwise       = case HM.keys m of
+        []     -> True
+        (k':_) -> isBottom $ HM.adjust (const bottom) k' m
+
 pInsertKeyStrict :: Int -> HashMap Key Int -> Bool
 pInsertKeyStrict v m = isBottom $ HM.insert bottom v m
 
@@ -130,11 +148,11 @@ pFromListWithValueResultStrict lst comb_lazy calc_good_raw
     calc_good Nothing y@(Just _) = cgr Nothing Nothing || cgr Nothing y
     calc_good x@(Just _) Nothing = cgr Nothing Nothing || cgr x Nothing
     calc_good x y = cgr Nothing Nothing || cgr Nothing y || cgr x Nothing || cgr x y
-    cgr = curry $ apply calc_good_raw
+    cgr = applyFun2 calc_good_raw
 
     -- The Maybe A -> Maybe A -> Maybe A that we're after, representing a
     -- potentially less total function than comb_lazy
-    comb x y = apply comb_lazy (x, y) <$ guard (calc_good x y)
+    comb x y = applyFun2 comb_lazy x y <$ guard (calc_good x y)
 
     -- What we get out of the conversion using fromListWith
     real_map = HM.fromListWith real_comb real_list
@@ -165,6 +183,8 @@ tests =
       , testProperty "delete is key-strict" $ keyStrict HM.delete
       , testProperty "adjust is key-strict" pAdjustKeyStrict
       , testProperty "adjust is value-strict" pAdjustValueStrict
+      , testProperty "adjustWithKey is key-strict" pAdjustWithKeyKeyStrict
+      , testProperty "adjustWithKey is value-strict" pAdjustWithKeyValueStrict
       , testProperty "insert is key-strict" pInsertKeyStrict
       , testProperty "insert is value-strict" pInsertValueStrict
       , testProperty "insertWith is key-strict" pInsertWithKeyStrict