From baf198701d9a1a2b0847ae7b23d4311097c9fbb0 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Wed, 4 Oct 2023 14:01:19 -0700
Subject: [PATCH] Updates: add getters, ctors and related widgets for erase

---
 .../open_addressing/open_addressing_impl.cuh  | 10 ++++++
 .../open_addressing_ref_impl.cuh              | 20 ++++++++++++
 include/cuco/detail/static_map/static_map.inl | 31 ++++++++++++++++---
 .../cuco/detail/static_map/static_map_ref.inl | 29 +++++++++++++++++
 include/cuco/detail/static_set/static_set.inl | 28 ++++++++++++++---
 .../cuco/detail/static_set/static_set_ref.inl | 22 +++++++++++++
 include/cuco/static_map.cuh                   |  7 +++++
 include/cuco/static_set.cuh                   |  7 +++++
 8 files changed, 145 insertions(+), 9 deletions(-)

diff --git a/include/cuco/detail/open_addressing/open_addressing_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_impl.cuh
index 8a24d8646..891feace7 100644
--- a/include/cuco/detail/open_addressing/open_addressing_impl.cuh
+++ b/include/cuco/detail/open_addressing/open_addressing_impl.cuh
@@ -639,6 +639,16 @@ class open_addressing_impl {
     return empty_key_sentinel_;
   }
 
+  /**
+   * @brief Gets the sentinel value used to represent an empty key slot.
+   *
+   * @return The sentinel value used to represent an empty key slot
+   */
+  [[nodiscard]] constexpr key_type erased_key_sentinel() const noexcept
+  {
+    return erased_key_sentinel_;
+  }
+
   /**
    * @brief Gets the key comparator.
    *
diff --git a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh
index 42db0976c..1133a4d53 100644
--- a/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh
+++ b/include/cuco/detail/open_addressing/open_addressing_ref_impl.cuh
@@ -127,6 +127,25 @@ class open_addressing_ref_impl {
   {
   }
 
+  /**
+   * @brief Constructs open_addressing_ref_impl.
+   *
+   * @param empty_slot_sentinel Sentinel indicating an empty slot
+   * @param probing_scheme Probing scheme
+   * @param storage_ref Non-owning ref of slot storage
+   */
+  __host__ __device__ explicit constexpr open_addressing_ref_impl(
+    value_type empty_slot_sentinel,
+    key_type erased_key_sentinel,
+    probing_scheme_type const& probing_scheme,
+    storage_ref_type storage_ref) noexcept
+    : empty_slot_sentinel_{empty_slot_sentinel},
+      erased_key_sentinel_{erased_key_sentinel},
+      probing_scheme_{probing_scheme},
+      storage_ref_{storage_ref}
+  {
+  }
+
   /**
    * @brief Gets the maximum number of elements the container can hold.
    *
@@ -920,6 +939,7 @@ class open_addressing_ref_impl {
   }
 
   value_type empty_slot_sentinel_;      ///< Sentinel value indicating an empty slot
+  key_type erased_key_sentinel_;        ///< Key value that represents an erased slot
   probing_scheme_type probing_scheme_;  ///< Probing scheme
   storage_ref_type storage_ref_;        ///< Slot storage ref
 };
diff --git a/include/cuco/detail/static_map/static_map.inl b/include/cuco/detail/static_map/static_map.inl
index 657e9f7eb..81cd49acc 100644
--- a/include/cuco/detail/static_map/static_map.inl
+++ b/include/cuco/detail/static_map/static_map.inl
@@ -468,6 +468,21 @@ constexpr static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator,
   return this->empty_value_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_type
+static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::
+  erased_key_sentinel() const noexcept
+{
+  return impl_->erased_key_sentinel();
+}
+
 template <class Key,
           class T,
           class Extent,
@@ -481,11 +496,17 @@ auto static_map<Key, T, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Stora
   Operators...) const noexcept
 {
   static_assert(sizeof...(Operators), "No operators specified");
-  return ref_type<Operators...>{cuco::empty_key<key_type>(this->empty_key_sentinel()),
-                                cuco::empty_value<mapped_type>(this->empty_value_sentinel()),
-                                impl_->key_eq(),
-                                impl_->probing_scheme(),
-                                impl_->storage_ref()};
+  return this->empty_key_sentinel() == this->erased_key_sentinel()
+           ? ref_type<Operators...>{cuco::empty_key<key_type>(this->empty_key_sentinel()),
+                                    cuco::empty_value<mapped_type>(this->empty_value_sentinel()),
+                                    impl_->key_eq(),
+                                    impl_->probing_scheme(),
+                                    impl_->storage_ref()}
+           : ref_type<Operators...>{cuco::empty_key<key_type>(this->empty_key_sentinel()),
+                                    cuco::empty_value<mapped_type>(this->empty_value_sentinel()),
+                                    impl_->key_eq(),
+                                    impl_->probing_scheme(),
+                                    impl_->storage_ref()};
 }
 }  // namespace experimental
 }  // namespace cuco
diff --git a/include/cuco/detail/static_map/static_map_ref.inl b/include/cuco/detail/static_map/static_map_ref.inl
index e85b77509..0789f6a62 100644
--- a/include/cuco/detail/static_map/static_map_ref.inl
+++ b/include/cuco/detail/static_map/static_map_ref.inl
@@ -50,6 +50,35 @@ __host__ __device__ constexpr static_map_ref<
 {
 }
 
+template <typename Key,
+          typename T,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr static_map_ref<
+  Key,
+  T,
+  Scope,
+  KeyEqual,
+  ProbingScheme,
+  StorageRef,
+  Operators...>::static_map_ref(cuco::empty_key<Key> empty_key_sentinel,
+                                cuco::empty_value<T> empty_value_sentinel,
+                                cuco::erased_key<Key> erased_key_sentinel,
+                                KeyEqual const& predicate,
+                                ProbingScheme const& probing_scheme,
+                                StorageRef storage_ref) noexcept
+  : impl_{cuco::pair{empty_key_sentinel, empty_value_sentinel},
+          erased_key_sentinel,
+          probing_scheme,
+          storage_ref},
+    empty_value_sentinel_{empty_value_sentinel},
+    predicate_{empty_key_sentinel, predicate}
+{
+}
+
 template <typename Key,
           typename T,
           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 6fe9d818f..09528122b 100644
--- a/include/cuco/detail/static_set/static_set.inl
+++ b/include/cuco/detail/static_set/static_set.inl
@@ -377,6 +377,20 @@ static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::emp
   return impl_->empty_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_type
+static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>::erased_key_sentinel()
+  const noexcept
+{
+  return impl_->erased_key_sentinel();
+}
+
 template <class Key,
           class Extent,
           cuda::thread_scope Scope,
@@ -389,10 +403,16 @@ auto static_set<Key, Extent, Scope, KeyEqual, ProbingScheme, Allocator, Storage>
   Operators...) const noexcept
 {
   static_assert(sizeof...(Operators), "No operators specified");
-  return ref_type<Operators...>{cuco::empty_key<key_type>(this->empty_key_sentinel()),
-                                impl_->key_eq(),
-                                impl_->probing_scheme(),
-                                impl_->storage_ref()};
+  return this->empty_key_sentinel() == this->erased_key_sentinel()
+           ? ref_type<Operators...>{cuco::empty_key<key_type>(this->empty_key_sentinel()),
+                                    impl_->key_eq(),
+                                    impl_->probing_scheme(),
+                                    impl_->storage_ref()}
+           : ref_type<Operators...>{cuco::empty_key<key_type>(this->empty_key_sentinel()),
+                                    cuco::erased_key<key_type>(this->erased_key_sentinel()),
+                                    impl_->key_eq(),
+                                    impl_->probing_scheme(),
+                                    impl_->storage_ref()};
 }
 }  // namespace experimental
 }  // namespace cuco
diff --git a/include/cuco/detail/static_set/static_set_ref.inl b/include/cuco/detail/static_set/static_set_ref.inl
index 3b754d972..69263f549 100644
--- a/include/cuco/detail/static_set/static_set_ref.inl
+++ b/include/cuco/detail/static_set/static_set_ref.inl
@@ -46,6 +46,28 @@ __host__ __device__ constexpr static_set_ref<
 {
 }
 
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr static_set_ref<
+  Key,
+  Scope,
+  KeyEqual,
+  ProbingScheme,
+  StorageRef,
+  Operators...>::static_set_ref(cuco::empty_key<Key> empty_key_sentinel,
+                                cuco::erased_key<Key> erased_key_sentinel,
+                                KeyEqual const& predicate,
+                                ProbingScheme const& probing_scheme,
+                                StorageRef storage_ref) noexcept
+  : impl_{empty_key_sentinel, erased_key_sentinel, probing_scheme, storage_ref},
+    predicate_{empty_key_sentinel, predicate}
+{
+}
+
 template <typename Key,
           cuda::thread_scope Scope,
           typename KeyEqual,
diff --git a/include/cuco/static_map.cuh b/include/cuco/static_map.cuh
index d541354fd..b3b56ca24 100644
--- a/include/cuco/static_map.cuh
+++ b/include/cuco/static_map.cuh
@@ -653,6 +653,13 @@ class static_map {
    */
   [[nodiscard]] constexpr mapped_type empty_value_sentinel() const noexcept;
 
+  /**
+   * @brief Gets the sentinel value used to represent an erased key slot.
+   *
+   * @return The sentinel value used to represent an erased key slot
+   */
+  [[nodiscard]] constexpr key_type erased_key_sentinel() const noexcept;
+
   /**
    * @brief Get device ref with operators.
    *
diff --git a/include/cuco/static_set.cuh b/include/cuco/static_set.cuh
index 252843e6d..e1a01c070 100644
--- a/include/cuco/static_set.cuh
+++ b/include/cuco/static_set.cuh
@@ -568,6 +568,13 @@ class static_set {
    */
   [[nodiscard]] constexpr key_type empty_key_sentinel() const noexcept;
 
+  /**
+   * @brief Gets the sentinel value used to represent an erased key slot.
+   *
+   * @return The sentinel value used to represent an erased key slot
+   */
+  [[nodiscard]] constexpr key_type erased_key_sentinel() const noexcept;
+
   /**
    * @brief Get device ref with operators.
    *