From 9fe6c82d08b95795e351af5dd5828fbb06a556c7 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Fri, 23 Aug 2024 16:22:08 -0700
Subject: [PATCH] Expose `key_eq` member accessors (#584)

Currently, `key_eq` is only accessible in ref classes, not in container
classes. This PR addresses the issue by exposing `key_eq` in all
classes.
---
 include/cuco/detail/static_map/static_map.inl    | 15 +++++++++++++++
 .../detail/static_multimap/static_multimap.inl   | 16 ++++++++++++++++
 .../detail/static_multiset/static_multiset.inl   | 15 +++++++++++++++
 include/cuco/detail/static_set/static_set.inl    | 13 +++++++++++++
 include/cuco/static_map.cuh                      |  7 +++++++
 include/cuco/static_multimap.cuh                 |  7 +++++++
 include/cuco/static_multiset.cuh                 |  7 +++++++
 include/cuco/static_set.cuh                      |  7 +++++++
 8 files changed, 87 insertions(+)

diff --git a/include/cuco/detail/static_map/static_map.inl b/include/cuco/detail/static_map/static_map.inl
index 08acdacaf..a90315187 100644
--- a/include/cuco/detail/static_map/static_map.inl
+++ b/include/cuco/detail/static_map/static_map.inl
@@ -709,6 +709,21 @@ static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
   return impl_->erased_key_sentinel();
 }
 
+template <class Key,
+          class T,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+constexpr static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::key_equal
+static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::key_eq()
+  const noexcept
+{
+  return impl_->key_eq();
+}
+
 template <class Key,
           class T,
           class Extent,
diff --git a/include/cuco/detail/static_multimap/static_multimap.inl b/include/cuco/detail/static_multimap/static_multimap.inl
index 3d97a2734..6582acd90 100644
--- a/include/cuco/detail/static_multimap/static_multimap.inl
+++ b/include/cuco/detail/static_multimap/static_multimap.inl
@@ -312,6 +312,22 @@ constexpr static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Alloca
   return impl_->erased_key_sentinel();
 }
 
+template <class Key,
+          class T,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+constexpr static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
+  key_equal
+  static_multimap<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::key_eq()
+    const noexcept
+{
+  return impl_->key_eq();
+}
+
 template <class Key,
           class T,
           class Extent,
diff --git a/include/cuco/detail/static_multiset/static_multiset.inl b/include/cuco/detail/static_multiset/static_multiset.inl
index b102540a3..b667ba5ae 100644
--- a/include/cuco/detail/static_multiset/static_multiset.inl
+++ b/include/cuco/detail/static_multiset/static_multiset.inl
@@ -393,6 +393,21 @@ static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>
   return impl_->erased_key_sentinel();
 }
 
+template <class Key,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+constexpr static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
+  key_equal
+  static_multiset<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::key_eq()
+    const noexcept
+{
+  return impl_->key_eq();
+}
+
 template <class Key,
           class Extent,
           cuda::thread_scope Scope,
diff --git a/include/cuco/detail/static_set/static_set.inl b/include/cuco/detail/static_set/static_set.inl
index bff86ab54..e8434d32d 100644
--- a/include/cuco/detail/static_set/static_set.inl
+++ b/include/cuco/detail/static_set/static_set.inl
@@ -514,6 +514,19 @@ static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::era
   return impl_->erased_key_sentinel();
 }
 
+template <class Key,
+          class Extent,
+          cuda::thread_scope Scope,
+          class KeyEqual,
+          class ProbingScheme,
+          class Allocator,
+          class Storage>
+constexpr static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::key_equal
+static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::key_eq() const noexcept
+{
+  return impl_->key_eq();
+}
+
 template <class Key,
           class Extent,
           cuda::thread_scope Scope,
diff --git a/include/cuco/static_map.cuh b/include/cuco/static_map.cuh
index 09e980261..899a0cdb4 100644
--- a/include/cuco/static_map.cuh
+++ b/include/cuco/static_map.cuh
@@ -952,6 +952,13 @@ class static_map {
    */
   [[nodiscard]] constexpr key_type erased_key_sentinel() const noexcept;
 
+  /**
+   * @brief Gets the function used to compare keys for equality
+   *
+   * @return The function used to compare keys for equality
+   */
+  [[nodiscard]] constexpr key_equal key_eq() const noexcept;
+
   /**
    * @brief Get device ref with operators.
    *
diff --git a/include/cuco/static_multimap.cuh b/include/cuco/static_multimap.cuh
index 31bac4b4f..043e04a9d 100644
--- a/include/cuco/static_multimap.cuh
+++ b/include/cuco/static_multimap.cuh
@@ -447,6 +447,13 @@ class static_multimap {
    */
   [[nodiscard]] constexpr key_type erased_key_sentinel() const noexcept;
 
+  /**
+   * @brief Gets the function used to compare keys for equality
+   *
+   * @return The function used to compare keys for equality
+   */
+  [[nodiscard]] constexpr key_equal key_eq() const noexcept;
+
   /**
    * @brief Get device ref with operators.
    *
diff --git a/include/cuco/static_multiset.cuh b/include/cuco/static_multiset.cuh
index 2158b31c4..55da49476 100644
--- a/include/cuco/static_multiset.cuh
+++ b/include/cuco/static_multiset.cuh
@@ -570,6 +570,13 @@ class static_multiset {
    */
   [[nodiscard]] constexpr key_type erased_key_sentinel() const noexcept;
 
+  /**
+   * @brief Gets the function used to compare keys for equality
+   *
+   * @return The function used to compare keys for equality
+   */
+  [[nodiscard]] constexpr key_equal key_eq() const noexcept;
+
   /**
    * @brief Get device ref with operators.
    *
diff --git a/include/cuco/static_set.cuh b/include/cuco/static_set.cuh
index 2ae20ed74..86c144f88 100644
--- a/include/cuco/static_set.cuh
+++ b/include/cuco/static_set.cuh
@@ -766,6 +766,13 @@ class static_set {
    */
   [[nodiscard]] constexpr key_type erased_key_sentinel() const noexcept;
 
+  /**
+   * @brief Gets the function used to compare keys for equality
+   *
+   * @return The function used to compare keys for equality
+   */
+  [[nodiscard]] constexpr key_equal key_eq() const noexcept;
+
   /**
    * @brief Get device ref with operators.
    *