From dcf85c09fe2254f13e680de1f29a008afeb85680 Mon Sep 17 00:00:00 2001
From: Martin Leitner-Ankerl <martin.ankerl@gmail.com>
Date: Sun, 16 Jul 2023 08:17:36 +0200
Subject: [PATCH] fix memory leak in move-assignment with different allocators

---
 include/ankerl/unordered_dense.h       | 27 ++++++++-----
 test/unit/pmr_move_with_allocators.cpp | 52 ++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 9 deletions(-)

diff --git a/include/ankerl/unordered_dense.h b/include/ankerl/unordered_dense.h
index 92fb49e7..6b320fbc 100644
--- 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
index d23cb5fa..f8f72788 100644
--- 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") {