From 1ba52414690c72f1a391c6318ee9c106f4d077aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Branimir=20Karad=C5=BEi=C4=87?= <branimirkaradzic@gmail.com>
Date: Wed, 26 Oct 2016 21:55:24 -0700
Subject: [PATCH] Cleanup.

---
 include/bx/handlealloc.h | 52 +++++++++++++++++++---------------------
 tests/handle_bench.cpp   | 25 +++++++++++++++++++
 2 files changed, 50 insertions(+), 27 deletions(-)

diff --git a/include/bx/handlealloc.h b/include/bx/handlealloc.h
index 856ae2d..bc8b42f 100644
--- a/include/bx/handlealloc.h
+++ b/include/bx/handlealloc.h
@@ -481,9 +481,7 @@ namespace bx
 			uint32_t idx = findIndex(_key);
 			if (UINT32_MAX != idx)
 			{
-				m_handle[idx] = invalid;
-				--m_numElements;
-				update();
+				removeIndex(idx);
 				return true;
 			}
 
@@ -494,21 +492,13 @@ namespace bx
 		{
 			if (invalid != _handle)
 			{
-				const uint32_t numElements = m_numElements;
 				for (uint32_t idx = 0; idx < MaxCapacityT; ++idx)
 				{
 					if (m_handle[idx] == _handle)
 					{
-						m_handle[idx] = invalid;
-						--m_numElements;
+						removeIndex(idx);
 					}
 				}
-
-				if (numElements != m_numElements)
-				{
-					update();
-					return true;
-				}
 			}
 
 			return false;
@@ -611,6 +601,29 @@ namespace bx
 			return UINT32_MAX;
 		}
 
+		void removeIndex(uint32_t _idx)
+		{
+			m_handle[_idx] = invalid;
+			--m_numElements;
+
+			for (uint32_t idx = (_idx + 1) % MaxCapacityT
+				; m_handle[idx] != invalid
+				; idx = (idx + 1) % MaxCapacityT)
+			{
+				if (m_handle[idx] != invalid)
+				{
+					const KeyT key = m_key[idx];
+					if (idx != findIndex(key) )
+					{
+						const uint16_t handle = m_handle[idx];
+						m_handle[idx] = invalid;
+						--m_numElements;
+						insert(key, handle);
+					}
+				}
+			}
+		}
+
 		uint32_t mix(uint32_t _x) const
 		{
 			const uint32_t tmp0   = uint32_mul(_x,   UINT32_C(2246822519) );
@@ -627,21 +640,6 @@ namespace bx
 			return result;
 		}
 
-		void update()
-		{
-			for (uint32_t idx = 0; idx < MaxCapacityT; ++idx)
-			{
-				if (m_handle[idx] != invalid)
-				{
-					const KeyT     key    = m_key[idx];
-					const uint16_t handle = m_handle[idx];
-					m_handle[idx] = invalid;
-					--m_numElements;
-					insert(key, handle);
-				}
-			}
-		}
-
 		uint32_t m_maxCapacity;
 		uint32_t m_numElements;
 
diff --git a/tests/handle_bench.cpp b/tests/handle_bench.cpp
index e230c4d..b34a659 100644
--- a/tests/handle_bench.cpp
+++ b/tests/handle_bench.cpp
@@ -1,6 +1,7 @@
 #include <bx/bx.h>
 #include <bx/timer.h>
 #include <bx/handlealloc.h>
+#include <bx/maputil.h>
 
 #include <tinystl/allocator.h>
 #include <tinystl/unordered_map.h>
@@ -28,6 +29,14 @@ int main()
 				tinystl::pair<TinyStlUnorderedMap::iterator, bool> ok = map.insert(tinystl::make_pair(uint64_t(jj), uint16_t(jj) ) );
 				assert(ok.second);
 			}
+
+			for (uint32_t jj = 0; jj < numElements; ++jj)
+			{
+				bool ok = bx::mapRemove(map, uint64_t(jj) );
+				assert(ok);
+			}
+
+			assert(map.size() == 0);
 		}
 
 		elapsed += bx::getHPCounter();
@@ -47,6 +56,14 @@ int main()
 				std::pair<StdUnorderedMap::iterator, bool> ok = map.insert(std::make_pair(uint64_t(jj), uint16_t(jj) ) );
 				assert(ok.second);
 			}
+
+			for (uint32_t jj = 0; jj < numElements; ++jj)
+			{
+				bool ok = bx::mapRemove(map, uint64_t(jj) );
+				assert(ok);
+			}
+
+			assert(map.size() == 0);
 		}
 
 		elapsed += bx::getHPCounter();
@@ -66,6 +83,14 @@ int main()
 				bool ok = map.insert(jj, uint16_t(jj) );
 				assert(ok);
 			}
+
+			for (uint32_t jj = 0; jj < numElements; ++jj)
+			{
+				bool ok = map.removeByKey(uint64_t(jj) );
+				assert(ok);
+			}
+
+			assert(map.getNumElements() == 0);
 		}
 
 		elapsed += bx::getHPCounter();