Fix leaks and crashes in OAHashMap

This changes the way the lifespan of items is managed to be consistent.

Bonus: Simplify cases of destroy-then-emplace.

2 files changed, 87 insertions, 16 deletions

diff --git a/core/oa_hash_map.h b/core/oa_hash_map.h
index e411ced044..4e2ea034ea 100644
--- a/core/oa_hash_map.h
+++ b/core/oa_hash_map.h
@@ -45,6 +45,9 @@
  *
  * The entries are stored inplace, so huge keys or values might fill cache lines
  * a lot faster.
+ *
+ * Only used keys and values are constructed. For free positions there's space
+ * in the arrays for each, but that memory is kept uninitialized.
  */
 template <class TKey, class TValue,
 		class Hasher = HashMapHasherDefault,
@@ -146,9 +149,9 @@ private:
 		uint32_t *old_hashes = hashes;
 
 		num_elements = 0;
-		keys = memnew_arr(TKey, capacity);
-		values = memnew_arr(TValue, capacity);
-		hashes = memnew_arr(uint32_t, capacity);
+		keys = static_cast<TKey *>(Memory::alloc_static(sizeof(TKey) * capacity));
+		values = static_cast<TValue *>(Memory::alloc_static(sizeof(TValue) * capacity));
+		hashes = static_cast<uint32_t *>(Memory::alloc_static(sizeof(uint32_t) * capacity));
 
 		for (uint32_t i = 0; i < capacity; i++) {
 			hashes[i] = 0;
@@ -160,11 +163,14 @@ private:
 			}
 
 			_insert_with_hash(old_hashes[i], old_keys[i], old_values[i]);
+
+			old_keys[i].~TKey();
+			old_values[i].~TValue();
 		}
 
-		memdelete_arr(old_keys);
-		memdelete_arr(old_values);
-		memdelete_arr(old_hashes);
+		Memory::free_static(old_keys);
+		Memory::free_static(old_values);
+		Memory::free_static(old_hashes);
 	}
 
 	void _resize_and_rehash() {
@@ -208,8 +214,7 @@ public:
 		bool exists = _lookup_pos(p_key, pos);
 
 		if (exists) {
-			values[pos].~TValue();
-			memnew_placement(&values[pos], TValue(p_data));
+			values[pos] = p_data;
 		} else {
 			insert(p_key, p_data);
 		}
@@ -226,8 +231,7 @@ public:
 		bool exists = _lookup_pos(p_key, pos);
 
 		if (exists) {
-			r_data.~TValue();
-			memnew_placement(&r_data, TValue(values[pos]));
+			r_data = values[pos];
 			return true;
 		}
 
@@ -343,9 +347,9 @@ public:
 	OAHashMap(uint32_t p_initial_capacity = 64) {
 		capacity = p_initial_capacity;
 
-		keys = memnew_arr(TKey, p_initial_capacity);
-		values = memnew_arr(TValue, p_initial_capacity);
-		hashes = memnew_arr(uint32_t, p_initial_capacity);
+		keys = static_cast<TKey *>(Memory::alloc_static(sizeof(TKey) * capacity));
+		values = static_cast<TValue *>(Memory::alloc_static(sizeof(TValue) * capacity));
+		hashes = static_cast<uint32_t *>(Memory::alloc_static(sizeof(uint32_t) * capacity));
 
 		for (uint32_t i = 0; i < p_initial_capacity; i++) {
 			hashes[i] = EMPTY_HASH;
@@ -353,9 +357,18 @@ public:
 	}
 
 	~OAHashMap() {
-		memdelete_arr(keys);
-		memdelete_arr(values);
-		memdelete_arr(hashes);
+		for (uint32_t i = 0; i < capacity; i++) {
+			if (hashes[i] == EMPTY_HASH) {
+				continue;
+			}
+
+			values[i].~TValue();
+			keys[i].~TKey();
+		}
+
+		Memory::free_static(keys);
+		Memory::free_static(values);
+		Memory::free_static(hashes);
 	}
 };
 
diff --git a/main/tests/test_oa_hash_map.cpp b/main/tests/test_oa_hash_map.cpp
index cae143bb5d..719817baf4 100644
--- a/main/tests/test_oa_hash_map.cpp
+++ b/main/tests/test_oa_hash_map.cpp
@@ -35,6 +35,37 @@
 
 namespace TestOAHashMap {
 
+struct CountedItem {
+	static int count;
+
+	int id = -1;
+	bool destroyed = false;
+
+	CountedItem() {
+		count++;
+	}
+
+	CountedItem(int p_id) :
+			id(p_id) {
+		count++;
+	}
+
+	CountedItem(const CountedItem &p_other) :
+			id(p_other.id) {
+		count++;
+	}
+
+	CountedItem &operator=(const CountedItem &p_other) = default;
+
+	~CountedItem() {
+		CRASH_COND(destroyed);
+		count--;
+		destroyed = true;
+	}
+};
+
+int CountedItem::count;
+
 MainLoop *test() {
 	OS::get_singleton()->print("\n\n\nHello from test\n");
 
@@ -152,6 +183,33 @@ MainLoop *test() {
 		map.set(5, 1);
 	}
 
+	// test memory management of items, should not crash or leak items
+	{
+		// Exercise different patterns of removal
+		for (int i = 0; i < 4; ++i) {
+			{
+				OAHashMap<String, CountedItem> map;
+				int id = 0;
+				for (int j = 0; j < 100; ++j) {
+					map.insert(itos(j), CountedItem(id));
+				}
+				if (i <= 1) {
+					for (int j = 0; j < 100; ++j) {
+						map.remove(itos(j));
+					}
+				}
+				if (i % 2 == 0) {
+					map.clear();
+				}
+			}
+
+			if (CountedItem::count != 0) {
+				OS::get_singleton()->print("%d != 0 (not performing the other test sub-cases, breaking...)\n", CountedItem::count);
+				break;
+			}
+		}
+	}
+
 	return nullptr;
 }