From fd68e36951fa25c42d11e17150362b71e76aa83d Mon Sep 17 00:00:00 2001
From: Simon Jakobi <simon.jakobi@gmail.com>
Date: Fri, 14 Nov 2025 00:18:59 +0100
Subject: [PATCH] Add HashMap.lookupKey and HashSet.lookupElement

Resolves #546.
---
 Data/HashMap/Internal.hs        | 35 +++++++++++++++++++++++++++++++++
 Data/HashMap/Internal/Strict.hs |  1 +
 Data/HashMap/Lazy.hs            |  1 +
 Data/HashMap/Strict.hs          |  1 +
 Data/HashSet.hs                 |  1 +
 Data/HashSet/Internal.hs        |  9 +++++++++
 tests/Properties/HashMapLazy.hs |  5 +++++
 7 files changed, 53 insertions(+)

diff --git a/Data/HashMap/Internal.hs b/Data/HashMap/Internal.hs
index ac5ed3ee..886ead2a 100644
--- a/Data/HashMap/Internal.hs
+++ b/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 =
diff --git a/Data/HashMap/Internal/Strict.hs b/Data/HashMap/Internal/Strict.hs
index d6218688..6951fa3b 100644
--- a/Data/HashMap/Internal/Strict.hs
+++ b/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
diff --git a/Data/HashMap/Lazy.hs b/Data/HashMap/Lazy.hs
index b0a01195..f981ca09 100644
--- a/Data/HashMap/Lazy.hs
+++ b/Data/HashMap/Lazy.hs
@@ -42,6 +42,7 @@ module Data.HashMap.Lazy
     , findWithDefault
     , lookupDefault
     , (!)
+    , lookupKey
     , insert
     , insertWith
     , delete
diff --git a/Data/HashMap/Strict.hs b/Data/HashMap/Strict.hs
index d0bc1e76..866e0b7a 100644
--- a/Data/HashMap/Strict.hs
+++ b/Data/HashMap/Strict.hs
@@ -41,6 +41,7 @@ module Data.HashMap.Strict
     , findWithDefault
     , lookupDefault
     , (!)
+    , lookupKey
     , insert
     , insertWith
     , delete
diff --git a/Data/HashSet.hs b/Data/HashSet.hs
index 330af38a..25b2d35c 100644
--- a/Data/HashSet.hs
+++ b/Data/HashSet.hs
@@ -107,6 +107,7 @@ module Data.HashSet
     , null
     , size
     , member
+    , lookupElement
     , insert
     , delete
     , isSubsetOf
diff --git a/Data/HashSet/Internal.hs b/Data/HashSet/Internal.hs
index 49736cb4..671c680f 100644
--- a/Data/HashSet/Internal.hs
+++ b/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
diff --git a/tests/Properties/HashMapLazy.hs b/tests/Properties/HashMapLazy.hs
index 1c556487..7f7f66e8 100644
--- a/tests/Properties/HashMapLazy.hs
+++ b/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 ->