From 115656e2b4a179660b2ff0e5f17c2ce8dba2decb Mon Sep 17 00:00:00 2001
From: Srinivas Yadav <43375352+srinivasyadav18@users.noreply.github.com>
Date: Wed, 7 Aug 2024 18:00:30 -0400
Subject: [PATCH] Expose `storage_ref` and `probing_scheme` in all container
 refs (#572)

This PR exposes `storage_ref` and `probing_scheme` in all container
refs.

---------

Co-authored-by: Yunsong Wang <yunsongw@nvidia.com>
---
 .../cuco/detail/static_map/static_map_ref.inl | 28 +++++++++++++++++++
 .../static_multimap/static_multimap_ref.inl   | 28 +++++++++++++++++++
 .../static_multiset/static_multiset_ref.inl   | 26 +++++++++++++++++
 .../cuco/detail/static_set/static_set_ref.inl | 26 +++++++++++++++++
 include/cuco/static_map_ref.cuh               | 14 ++++++++++
 include/cuco/static_multimap_ref.cuh          | 14 ++++++++++
 include/cuco/static_multiset_ref.cuh          | 14 ++++++++++
 include/cuco/static_set_ref.cuh               | 14 ++++++++++
 8 files changed, 164 insertions(+)

diff --git a/include/cuco/detail/static_map/static_map_ref.inl b/include/cuco/detail/static_map/static_map_ref.inl
index 39ef375bc..b38562027 100644
--- a/include/cuco/detail/static_map/static_map_ref.inl
+++ b/include/cuco/detail/static_map/static_map_ref.inl
@@ -178,6 +178,34 @@ static_map_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>
   return impl_.capacity();
 }
 
+template <typename Key,
+          typename T,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr auto
+static_map_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::storage_ref()
+  const noexcept
+{
+  return this->impl_.storage_ref();
+}
+
+template <typename Key,
+          typename T,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr auto
+static_map_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::probing_scheme()
+  const noexcept
+{
+  return this->impl_.probing_scheme();
+}
+
 template <typename Key,
           typename T,
           cuda::thread_scope Scope,
diff --git a/include/cuco/detail/static_multimap/static_multimap_ref.inl b/include/cuco/detail/static_multimap/static_multimap_ref.inl
index a79edba42..01f9171d4 100644
--- a/include/cuco/detail/static_multimap/static_multimap_ref.inl
+++ b/include/cuco/detail/static_multimap/static_multimap_ref.inl
@@ -176,6 +176,34 @@ static_multimap_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operator
   return impl_.capacity();
 }
 
+template <typename Key,
+          typename T,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr auto
+static_multimap_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::storage_ref()
+  const noexcept
+{
+  return this->impl_.storage_ref();
+}
+
+template <typename Key,
+          typename T,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr auto
+static_multimap_ref<Key, T, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::
+  probing_scheme() const noexcept
+{
+  return this->impl_.probing_scheme();
+}
+
 template <typename Key,
           typename T,
           cuda::thread_scope Scope,
diff --git a/include/cuco/detail/static_multiset/static_multiset_ref.inl b/include/cuco/detail/static_multiset/static_multiset_ref.inl
index d08e50f0e..fc1f3db9d 100644
--- a/include/cuco/detail/static_multiset/static_multiset_ref.inl
+++ b/include/cuco/detail/static_multiset/static_multiset_ref.inl
@@ -121,6 +121,32 @@ static_multiset_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators..
   return impl_.capacity();
 }
 
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr auto
+static_multiset_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::storage_ref()
+  const noexcept
+{
+  return this->impl_.storage_ref();
+}
+
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr auto
+static_multiset_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::probing_scheme()
+  const noexcept
+{
+  return this->impl_.probing_scheme();
+}
+
 template <typename Key,
           cuda::thread_scope Scope,
           typename KeyEqual,
diff --git a/include/cuco/detail/static_set/static_set_ref.inl b/include/cuco/detail/static_set/static_set_ref.inl
index ebec75672..dc25c2372 100644
--- a/include/cuco/detail/static_set/static_set_ref.inl
+++ b/include/cuco/detail/static_set/static_set_ref.inl
@@ -153,6 +153,32 @@ static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::c
   return impl_.capacity();
 }
 
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr auto
+static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::storage_ref()
+  const noexcept
+{
+  return this->impl_.storage_ref();
+}
+
+template <typename Key,
+          cuda::thread_scope Scope,
+          typename KeyEqual,
+          typename ProbingScheme,
+          typename StorageRef,
+          typename... Operators>
+__host__ __device__ constexpr auto
+static_set_ref<Key, Scope, KeyEqual, ProbingScheme, StorageRef, Operators...>::probing_scheme()
+  const noexcept
+{
+  return this->impl_.probing_scheme();
+}
+
 template <typename Key,
           cuda::thread_scope Scope,
           typename KeyEqual,
diff --git a/include/cuco/static_map_ref.cuh b/include/cuco/static_map_ref.cuh
index 4006d84ae..e12bdb6f2 100644
--- a/include/cuco/static_map_ref.cuh
+++ b/include/cuco/static_map_ref.cuh
@@ -204,6 +204,20 @@ class static_map_ref
    */
   [[nodiscard]] __host__ __device__ constexpr iterator end() noexcept;
 
+  /**
+   * @brief Gets the non-owning storage ref.
+   *
+   * @return The non-owning storage ref of the container
+   */
+  [[nodiscard]] __host__ __device__ constexpr auto storage_ref() const noexcept;
+
+  /**
+   * @brief Gets the probing scheme.
+   *
+   * @return The probing scheme used for the container
+   */
+  [[nodiscard]] __host__ __device__ constexpr auto probing_scheme() const noexcept;
+
   /**
    * @brief Creates a reference with new operators from the current object.
    *
diff --git a/include/cuco/static_multimap_ref.cuh b/include/cuco/static_multimap_ref.cuh
index fceb10489..559392db1 100644
--- a/include/cuco/static_multimap_ref.cuh
+++ b/include/cuco/static_multimap_ref.cuh
@@ -206,6 +206,20 @@ class static_multimap_ref
    */
   [[nodiscard]] __host__ __device__ constexpr iterator end() noexcept;
 
+  /**
+   * @brief Gets the non-owning storage ref.
+   *
+   * @return The non-owning storage ref of the container
+   */
+  [[nodiscard]] __host__ __device__ constexpr auto storage_ref() const noexcept;
+
+  /**
+   * @brief Gets the probing scheme.
+   *
+   * @return The probing scheme used for the container
+   */
+  [[nodiscard]] __host__ __device__ constexpr auto probing_scheme() const noexcept;
+
   /**
    * @brief Creates a reference with new operators from the current object.
    *
diff --git a/include/cuco/static_multiset_ref.cuh b/include/cuco/static_multiset_ref.cuh
index dc352ea02..9d9ed5307 100644
--- a/include/cuco/static_multiset_ref.cuh
+++ b/include/cuco/static_multiset_ref.cuh
@@ -183,6 +183,20 @@ class static_multiset_ref
    */
   [[nodiscard]] __host__ __device__ constexpr iterator end() noexcept;
 
+  /**
+   * @brief Gets the non-owning storage ref.
+   *
+   * @return The non-owning storage ref of the container
+   */
+  [[nodiscard]] __host__ __device__ constexpr auto storage_ref() const noexcept;
+
+  /**
+   * @brief Gets the probing scheme.
+   *
+   * @return The probing scheme used for the container
+   */
+  [[nodiscard]] __host__ __device__ constexpr auto probing_scheme() const noexcept;
+
   /**
    * @brief Creates a reference with new operators from the current object.
    *
diff --git a/include/cuco/static_set_ref.cuh b/include/cuco/static_set_ref.cuh
index 2519f6609..a7c691b0e 100644
--- a/include/cuco/static_set_ref.cuh
+++ b/include/cuco/static_set_ref.cuh
@@ -181,6 +181,20 @@ class static_set_ref
    */
   [[nodiscard]] __host__ __device__ constexpr iterator end() noexcept;
 
+  /**
+   * @brief Gets the non-owning storage ref.
+   *
+   * @return The non-owning storage ref of the container
+   */
+  [[nodiscard]] __host__ __device__ constexpr auto storage_ref() const noexcept;
+
+  /**
+   * @brief Gets the probing scheme.
+   *
+   * @return The probing scheme used for the container
+   */
+  [[nodiscard]] __host__ __device__ constexpr auto probing_scheme() const noexcept;
+
   /**
    * @brief Creates a reference with new operators from the current object.
    *