Adding alter function, documentation, and tests.

Author: gibiansky

Data/HashMap/Base.hs

@@ -26,6 +26,7 @@ module Data.HashMap.Base
     , unsafeInsert
     , delete
     , adjust
+    , alter
 
       -- * Combine
       -- ** Union
@@ -574,6 +575,16 @@ adjust f k0 m0 = go h0 k0 0 m0
         | otherwise = t
 {-# INLINABLE adjust #-}
 
+-- | /O(log n)/  The expression (@'alter' f k map@) alters the value @x@ at @k@, or
+-- absence thereof. @alter@ can be used to insert, delete, or update a value in a
+-- map. In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
+alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
+alter f k m =
+  case f (lookup k m) of
+    Nothing -> delete k m
+    Just v  -> insert k v m
+{-# INLINABLE alter #-}
+
 ------------------------------------------------------------------------
 -- * Combine
 

Data/HashMap/Lazy.hs

@@ -47,6 +47,7 @@ module Data.HashMap.Lazy
     , insertWith
     , delete
     , adjust
+    , alter
 
       -- * Combine
       -- ** Union

Data/HashMap/Strict.hs

@@ -47,6 +47,7 @@ module Data.HashMap.Strict
     , insertWith
     , delete
     , adjust
+    , alter
 
       -- * Combine
       -- ** Union
@@ -92,7 +93,7 @@ import Prelude hiding (map)
 import qualified Data.HashMap.Array as A
 import qualified Data.HashMap.Base as HM
 import Data.HashMap.Base hiding (
-    adjust, fromList, fromListWith, insert, insertWith, intersectionWith,
+    alter, adjust, fromList, fromListWith, insert, insertWith, intersectionWith,
     map, mapWithKey, singleton, unionWith)
 import Data.HashMap.Unsafe (runST)
 
@@ -227,6 +228,16 @@ adjust f k0 m0 = go h0 k0 0 m0
         | otherwise = t
 {-# INLINABLE adjust #-}
 
+-- | /O(log n)/  The expression (@'alter' f k map@) alters the value @x@ at @k@, or
+-- absence thereof. @alter@ can be used to insert, delete, or update a value in a
+-- map. In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@.
+alter :: (Eq k, Hashable k) => (Maybe v -> Maybe v) -> k -> HashMap k v -> HashMap k v
+alter f k m =
+  case f (HM.lookup k m) of
+    Nothing -> delete k m
+    Just v  -> insert k v m
+{-# INLINABLE alter #-}
+
 ------------------------------------------------------------------------
 -- * Combine
 

tests/HashMapProperties.hs

@@ -98,6 +98,15 @@ pInsertWith k = M.insertWith (+) k 1 `eq_` HM.insertWith (+) k 1
 pAdjust :: Key -> [(Key, Int)] -> Bool
 pAdjust k = M.adjust succ k `eq_` HM.adjust succ k
 
+pAlterAdjust :: Key -> [(Key, Int)] -> Bool
+pAlterAdjust k = M.alter (fmap succ) k `eq_` HM.alter (fmap succ) k
+
+pAlterInsert :: Key -> [(Key, Int)] -> Bool
+pAlterInsert k = M.alter (const $ Just 3) k `eq_` HM.alter (const $ Just 3) k
+
+pAlterDelete :: Key -> [(Key, Int)] -> Bool
+pAlterDelete k = M.alter (const Nothing) k `eq_` HM.alter (const Nothing) k
+
 ------------------------------------------------------------------------
 -- ** Combine
 
@@ -209,6 +218,9 @@ tests =
       , testProperty "deleteCollision" pDeleteCollision
       , testProperty "insertWith" pInsertWith
       , testProperty "adjust" pAdjust
+      , testProperty "alterAdjust" pAlterAdjust
+      , testProperty "alterInsert" pAlterInsert
+      , testProperty "alterDelete" pAlterDelete
       ]
     -- Combine
     , testProperty "union" pUnion