Optimize alter (#404)

Author: sjakobi

Data/HashMap/Internal.hs

@@ -134,6 +134,7 @@ module Data.HashMap.Internal
     , equalKeys1
     , lookupRecordCollision
     , LookupRes(..)
+    , lookupResToMaybe
     , insert'
     , delete'
     , lookup'
@@ -655,6 +656,11 @@ lookup' h k m = case lookupRecordCollision# h k m of
 -- If a collision did not occur then it will have the Int value (-1).
 data LookupRes a = Absent | Present a !Int
 
+lookupResToMaybe :: LookupRes a -> Maybe a
+lookupResToMaybe Absent        = Nothing
+lookupResToMaybe (Present x _) = Just x
+{-# INLINE lookupResToMaybe #-}
+
 -- Internal helper for lookup. This version takes the precomputed hash so
 -- that functions that make multiple calls to lookup and related functions
 -- (insert, delete) only need to calculate the hash once.
@@ -1265,11 +1271,19 @@ update f = alter (>>= f)
 -- '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
--- TODO(m-renaud): Consider using specialized insert and delete for alter.
 alter f k m =
-  case f (lookup k m) of
-    Nothing -> delete k m
-    Just v  -> insert k v m
+    let !h = hash k
+        !lookupRes = lookupRecordCollision h k m
+    in case f (lookupResToMaybe lookupRes) of
+        Nothing -> case lookupRes of
+            Absent            -> m
+            Present _ collPos -> deleteKeyExists collPos h k m
+        Just v' -> case lookupRes of
+            Absent            -> insertNewKey h k v' m
+            Present v collPos ->
+                if v `ptrEq` v'
+                    then m
+                    else insertKeyExists collPos h k v' m
 {-# INLINABLE alter #-}
 
 -- | \(O(\log n)\)  The expression @('alterF' f k map)@ alters the value @x@ at

Data/HashMap/Internal/Strict.hs

@@ -306,9 +306,18 @@ update f = alter (>>= f)
 -- @
 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 -> HM.delete k m
-    Just v  -> insert k v m
+    let !h = hash k
+        !lookupRes = HM.lookupRecordCollision h k m
+    in case f (HM.lookupResToMaybe lookupRes) of
+        Nothing -> case lookupRes of
+            Absent            -> m
+            Present _ collPos -> HM.deleteKeyExists collPos h k m
+        Just !v' -> case lookupRes of
+            Absent             -> HM.insertNewKey h k v' m
+            Present v collPos ->
+                if v `ptrEq` v'
+                    then m
+                    else HM.insertKeyExists collPos h k v' m
 {-# INLINABLE alter #-}
 
 -- | \(O(\log n)\)  The expression (@'alterF' f k map@) alters the value @x@ at