fix memory leak in move-assignment with different allocators

martinus · Jul 16, 2023 · dcf85c0 · dcf85c0
1 parent 135b11e
commit dcf85c0
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 9 deletions.
diff --git a/include/ankerl/unordered_dense.h b/include/ankerl/unordered_dense.h
@@ -1214,21 +1214,30 @@ class table : public std::conditional_t<is_map_v<T>, base_table_type_map<T>, bas
         if (&other != this) {
             deallocate_buckets(); // deallocate before m_values is set (might have another allocator)
             m_values = std::move(other.m_values);
+            other.m_values.clear();
 
-            // we can only reuse m_buckets over when both maps have the same allocator!
+            // we can only reuse m_buckets when both maps have the same allocator!
             if (get_allocator() == other.get_allocator()) {
                 m_buckets = std::exchange(other.m_buckets, nullptr);
+                m_num_buckets = std::exchange(other.m_num_buckets, 0);
+                m_max_bucket_capacity = std::exchange(other.m_max_bucket_capacity, 0);
+                m_shifts = std::exchange(other.m_shifts, initial_shifts);
+                m_max_load_factor = std::exchange(other.m_max_load_factor, default_max_load_factor);
+                m_hash = std::exchange(other.m_hash, {});
+                m_equal = std::exchange(other.m_equal, {});
             } else {
+                // set max_load_factor *before* copying the other's buckets, so we have the same
+                // behavior
+                m_max_load_factor = other.m_max_load_factor;
+
+                // copy_buckets sets m_buckets, m_num_buckets, m_max_bucket_capacity, m_shifts
                 copy_buckets(other);
+                // clear's the other's buckets so other is now already usable. 
+                other.clear_buckets();
+                m_hash = other.m_hash;
+                m_equal = other.m_equal;
             }
-
-            m_num_buckets = std::exchange(other.m_num_buckets, 0);
-            m_max_bucket_capacity = std::exchange(other.m_max_bucket_capacity, 0);
-            m_max_load_factor = std::exchange(other.m_max_load_factor, default_max_load_factor);
-            m_hash = std::exchange(other.m_hash, {});
-            m_equal = std::exchange(other.m_equal, {});
-            m_shifts = std::exchange(other.m_shifts, initial_shifts);
-            other.m_values.clear();
+            // map "other" is now already usable, it's empty.
         }
         return *this;
     }

diff --git a/test/unit/pmr_move_with_allocators.cpp b/test/unit/pmr_move_with_allocators.cpp
@@ -40,6 +40,58 @@ auto doTest() {
         Map m(return_hello_world<Map>(&pool), ANKERL_UNORDERED_DENSE_PMR::new_delete_resource());
         REQUIRE(m.contains(0));
     }
+
+    {
+        Map a(ANKERL_UNORDERED_DENSE_PMR::new_delete_resource());
+        a[0] = "hello";
+        Map b(&pool);
+        b[0] = "world";
+
+        // looping here causes lots of memory held up
+        // in the resources
+        for (int i = 0; i < 100; ++i) {
+            std::swap(a, b);
+            REQUIRE(b[0] == "hello");
+            REQUIRE(a[0] == "world");
+
+            std::swap(a, b);
+            REQUIRE(a[0] == "hello");
+            REQUIRE(b[0] == "world");
+        }
+    }
+
+    {
+        Map a(ANKERL_UNORDERED_DENSE_PMR::new_delete_resource());
+        a[0] = "hello";
+        Map b(&pool);
+        b[0] = "world";
+
+        // looping here causes lots of memory held up
+        // in the resources
+        for (int i = 0; i < 100; ++i) {
+            std::swap(a, b);
+            REQUIRE(b[0] == "hello");
+            REQUIRE(a[0] == "world");
+
+            std::swap(a, b);
+            REQUIRE(a[0] == "hello");
+            REQUIRE(b[0] == "world");
+        }
+    }
+
+    {
+        Map a(&pool);
+        a[0] = "world";
+
+        Map tmp(ANKERL_UNORDERED_DENSE_PMR::new_delete_resource());
+        tmp[0] = "nope";
+        tmp = std::move(a);
+        REQUIRE(tmp[0] == "world");
+        REQUIRE(a.empty());
+        a[0] = "hey";
+        REQUIRE(a.size() == 1);
+        REQUIRE(a[0] == "hey");
+    }
 }
 
 TEST_CASE("move_with_allocators") {