From 8349635ffcd85f23e43bf89e9559a76b511d0e5b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pedro=20J=2E=20Est=C3=A9banez?= <pedrojrulez@gmail.com>
Date: Fri, 15 May 2020 00:48:04 +0200
Subject: [PATCH] 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.
---
 core/oa_hash_map.h              | 46 +++++++++++++++---------
 main/tests/test_oa_hash_map.cpp | 62 +++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 16 deletions(-)

diff --git a/core/oa_hash_map.h b/core/oa_hash_map.h
index 7407c528164..4770e0ec454 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,
@@ -149,9 +152,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;
@@ -163,11 +166,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() {
@@ -214,8 +220,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);
 		}
@@ -232,8 +237,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;
 		}
 
@@ -336,9 +340,9 @@ public:
 		capacity = p_initial_capacity;
 		num_elements = 0;
 
-		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;
@@ -347,9 +351,19 @@ 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 ac53f124d2f..8623a72b586 100644
--- a/main/tests/test_oa_hash_map.cpp
+++ b/main/tests/test_oa_hash_map.cpp
@@ -35,6 +35,41 @@
 
 namespace TestOAHashMap {
 
+struct CountedItem {
+	static int count;
+
+	int id;
+	bool destroyed;
+
+	CountedItem() :
+			id(-1),
+			destroyed(false) {
+		count++;
+	}
+
+	CountedItem(int p_id) :
+			id(p_id),
+			destroyed(false) {
+		count++;
+	}
+
+	CountedItem(const CountedItem &p_other) :
+			id(p_other.id),
+			destroyed(false) {
+		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 +187,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 NULL;
 }
 } // namespace TestOAHashMap