Optimize Strict.alter

Author: sjakobi

Data/HashMap/Internal.hs

@@ -128,6 +128,7 @@ module Data.HashMap.Internal
     , equalKeys1
     , lookupRecordCollision
     , LookupRes(..)
+    , lookupResToMaybe
     , insert'
     , delete'
     , lookup'
@@ -613,6 +614,10 @@ 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
+
 -- 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.

Data/HashMap/Internal/Strict.hs

@@ -305,9 +305,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