Add HashMap.lookupKey and HashSet.lookupElement (#554)

Resolves #546.

Author: sjakobi

Data/HashMap/Internal.hs

@@ -47,6 +47,7 @@ module Data.HashMap.Internal
     , findWithDefault
     , lookupDefault
     , (!)
+    , lookupKey
     , insert
     , insertWith
     , unsafeInsert
@@ -784,6 +785,40 @@ lookupDefault = findWithDefault
 
 infixl 9 !
 
+-- | \(O(\log n)\) For a given key, return the equal key stored in the map,
+-- if present, otherwise return 'Nothing'.
+--
+-- This function can be used for /interning/, i.e. to reduce memory usage.
+lookupKey :: Hashable k => k -> HashMap k v -> Maybe k
+lookupKey k = \m -> fromMaybe# (lookupKeyInSubtree# 0 (hash k) k m)
+  where
+    fromMaybe# (# (##) | #) = Nothing
+    fromMaybe# (# | a #) = Just a
+{-# INLINE lookupKey #-}
+
+lookupKeyInSubtree# :: Eq k => Shift -> Hash -> k -> HashMap k v -> (# (##) | k #)
+lookupKeyInSubtree# !s !hx kx = \case
+  Empty -> (# (##) | #)
+  Leaf hy (L ky _)
+    | hx == hy && kx == ky -> (# | ky #)
+    | otherwise -> (# (##) | #)
+  BitmapIndexed b ary
+    | m .&. b == 0 -> (# (##) | #)
+    | otherwise -> case A.index# ary i of
+        (# st #) -> lookupKeyInSubtree# (nextShift s) hx kx st
+    where
+      m = mask hx s
+      i = sparseIndex b m
+  Full ary -> case A.index# ary (index hx s) of
+    (# st #) -> lookupKeyInSubtree# (nextShift s) hx kx st
+  Collision hy ary
+    | hx == hy
+    , Just i <- indexOf kx ary
+    , (# L ky _ #) <- A.index# ary i
+    -> (# | ky #)
+    | otherwise -> (# (##) | #)
+{-# INLINABLE lookupKeyInSubtree# #-}
+
 -- | Create a 'Collision' value with two 'Leaf' values.
 collision :: Hash -> Leaf k v -> Leaf k v -> HashMap k v
 collision h !e1 !e2 =

Data/HashMap/Internal/Strict.hs

@@ -62,6 +62,7 @@ module Data.HashMap.Internal.Strict
     , HM.findWithDefault
     , HM.lookupDefault
     , (HM.!)
+    , HM.lookupKey
     , insert
     , insertWith
     , HM.delete

Data/HashMap/Lazy.hs

@@ -42,6 +42,7 @@ module Data.HashMap.Lazy
     , findWithDefault
     , lookupDefault
     , (!)
+    , lookupKey
     , insert
     , insertWith
     , delete

Data/HashMap/Strict.hs

@@ -41,6 +41,7 @@ module Data.HashMap.Strict
     , findWithDefault
     , lookupDefault
     , (!)
+    , lookupKey
     , insert
     , insertWith
     , delete

Data/HashSet.hs

@@ -107,6 +107,7 @@ module Data.HashSet
     , null
     , size
     , member
+    , lookupElement
     , insert
     , delete
     , isSubsetOf

Data/HashSet/Internal.hs

@@ -52,6 +52,7 @@ module Data.HashSet.Internal
     , null
     , size
     , member
+    , lookupElement
     , insert
     , delete
     , isSubsetOf
@@ -350,6 +351,14 @@ member a s = case H.lookup a (asMap s) of
                _      -> False
 {-# INLINABLE member #-}
 
+-- | \(O(\log n)\) For a given value, return the equal element in the set if
+-- present, otherwise return 'Nothing'.
+--
+-- This is useful for /interning/, i.e. to reduce memory usage.
+lookupElement :: Hashable a => a -> HashSet a -> Maybe a
+lookupElement a = H.lookupKey a . asMap
+{-# INLINE lookupElement #-}
+
 -- | \(O(\log n)\) Add the specified value to this set.
 --
 -- >>> HashSet.insert 1 HashSet.empty

tests/Properties/HashMapLazy.hs

@@ -24,6 +24,7 @@ import Data.Function               (on)
 import Data.Functor.Identity       (Identity (..))
 import Data.Hashable               (Hashable (hashWithSalt))
 import Data.HashMap.Internal.Debug (Validity (..), valid)
+import Data.Maybe                  (isJust)
 import Data.Ord                    (comparing)
 import Test.QuickCheck             (Arbitrary (..), Fun, Property, pattern Fn,
                                     pattern Fn2, pattern Fn3, (===), (==>))
@@ -173,6 +174,10 @@ tests =
       \(k :: Key) (m :: HMKI) -> HM.lookup k m === M.lookup k (toOrdMap m)
     , testProperty "!?" $
       \(k :: Key) (m :: HMKI) -> m HM.!? k === M.lookup k (toOrdMap m)
+    , testGroup "lookupKey" $
+      [ testProperty "isJust (lookupKey k m) == member k m" $
+        \(k :: Key) (m :: HMKI) -> isJust (HM.lookupKey k m) === HM.member k m
+      ]
     , testGroup "insert"
       [ testProperty "model" $
         \(k :: Key) (v :: Int) x ->