Enable rehash tests with erase

NVIDIA · Oct 11, 2023 · e3c6fd2 · e3c6fd2
1 parent 037baf2
commit e3c6fd2
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git 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
@@ -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);
 }