From e3c6fd2c8653d9858b4be3dc842df90807e55e26 Mon Sep 17 00:00:00 2001
From: Daniel Juenger <2955913+sleeepyjack@users.noreply.github.com>
Date: Wed, 11 Oct 2023 13:05:27 +0000
Subject: [PATCH] Enable rehash tests with erase

---
 tests/static_map/rehash_test.cu | 14 +++++++++-----
 tests/static_set/rehash_test.cu | 14 +++++++++-----
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/tests/static_map/rehash_test.cu b/tests/static_map/rehash_test.cu
index 11ed38534..69a73c6b3 100644
--- a/tests/static_map/rehash_test.cu
+++ b/tests/static_map/rehash_test.cu
@@ -27,10 +27,14 @@ TEST_CASE("Rehash", "")
 {
   using key_type    = int;
   using mapped_type = long;
+
   constexpr std::size_t num_keys{400};
+  constexpr std::size_t num_erased_keys{100};
 
-  cuco::experimental::static_map<key_type, mapped_type> map{
-    num_keys, cuco::empty_key<key_type>{-1}, cuco::empty_value<mapped_type>{-1}};
+  cuco::experimental::static_map map{num_keys,
+                                     cuco::empty_key<key_type>{-1},
+                                     cuco::empty_value<mapped_type>{-1},
+                                     cuco::erased_key<key_type>{-2}};
 
   thrust::device_vector<key_type> d_keys(num_keys);
   thrust::device_vector<mapped_type> d_values(num_keys);
@@ -49,7 +53,7 @@ TEST_CASE("Rehash", "")
   map.rehash(num_keys * 2);
   REQUIRE(map.size() == num_keys);
 
-  // TODO erase num_erased keys
-  // map.rehash()
-  // REQUIRE(map.size() == num_keys - num_erased);
+  map.erase(d_keys.begin(), d_keys.begin() + num_erased_keys);
+  map.rehash();
+  REQUIRE(map.size() == num_keys - num_erased_keys);
 }
diff --git a/tests/static_set/rehash_test.cu b/tests/static_set/rehash_test.cu
index b2b2bc9ca..98106df8a 100644
--- a/tests/static_set/rehash_test.cu
+++ b/tests/static_set/rehash_test.cu
@@ -23,11 +23,15 @@
 
 TEST_CASE("Rehash", "")
 {
+  using key_type = int;
+
   constexpr std::size_t num_keys{400};
+  constexpr std::size_t num_erased_keys{100};
 
-  cuco::experimental::static_set<int> set{num_keys, cuco::empty_key{-1}};
+  cuco::experimental::static_set set{
+    num_keys, cuco::empty_key<key_type>{-1}, cuco::erased_key<key_type>{-2}};
 
-  thrust::device_vector<int> d_keys(num_keys);
+  thrust::device_vector<key_type> d_keys(num_keys);
 
   thrust::sequence(d_keys.begin(), d_keys.end());
 
@@ -39,7 +43,7 @@ TEST_CASE("Rehash", "")
   set.rehash(num_keys * 2);
   REQUIRE(set.size() == num_keys);
 
-  // TODO erase num_erased keys
-  // set.rehash()
-  // REQUIRE(set.size() == num_keys - num_erased);
+  set.erase(d_keys.begin(), d_keys.begin() + num_erased_keys);
+  set.rehash();
+  REQUIRE(set.size() == num_keys - num_erased_keys);
 }