Core/MemoryPool: Add iterator

include/Nazara/Core/MemoryPool.hpp

@@ -9,6 +9,7 @@
 
 #include <Nazara/Prerequisites.hpp>
 #include <Nazara/Core/Bitset.hpp>
+#include <iterator>
 #include <memory>
 #include <vector>
 
@@ -18,6 +19,9 @@ namespace Nz
 	class MemoryPool
 	{
 		public:
+			class iterator;
+			friend iterator;
+
 			MemoryPool(std::size_t blockSize);
 			MemoryPool(const MemoryPool&) = delete;
 			MemoryPool(MemoryPool&&) noexcept = default;
@@ -36,8 +40,14 @@ namespace Nz
 
 			void Reset();
 
+			T* RetrieveFromIndex(std::size_t index);
 			std::size_t RetrieveEntryIndex(const T* data);
 
+			// std interface
+			iterator begin();
+			iterator end();
+			std::size_t size();
+
 			MemoryPool& operator=(const MemoryPool&) = delete;
 			MemoryPool& operator=(MemoryPool&& pool) noexcept = default;
 
@@ -45,6 +55,10 @@ namespace Nz
 
 		private:
 			void AllocateBlock();
+			T* GetAllocatedPointer(std::size_t blockIndex, std::size_t localIndex);
+			std::pair<std::size_t, std::size_t> GetFirstAllocatedEntry() const;
+			std::pair<std::size_t, std::size_t> GetFirstAllocatedEntryFromBlock(std::size_t blockIndex) const;
+			std::pair<std::size_t, std::size_t> GetNextAllocatedEntry(std::size_t blockIndex, std::size_t localIndex) const;
 
 			using AlignedStorage = std::aligned_storage_t<sizeof(T), Alignment>;
 
@@ -53,11 +67,45 @@ namespace Nz
 				std::size_t occupiedEntryCount = 0;
 				std::unique_ptr<AlignedStorage[]> memory;
 				Bitset<UInt64> freeEntries;
+				Bitset<UInt64> occupiedEntries; //< Opposite of freeEntries
 			};
 
 			std::size_t m_blockSize;
 			std::vector<Block> m_blocks;
 	};
+
+	template<typename T, std::size_t Alignment>
+	class MemoryPool<T, Alignment>::iterator
+	{
+		friend MemoryPool;
+
+		public:
+			using iterator_category = std::input_iterator_tag;
+			using value_type = T;
+			using difference_type = std::ptrdiff_t;
+			using pointer = T*;
+			using reference = T&;
+
+			iterator(const iterator&) = default;
+			iterator(iterator&&) = default;
+
+			iterator& operator=(const iterator&) = default;
+			iterator& operator=(iterator&&) = default;
+
+			iterator operator++(int);
+			iterator& operator++();
+
+			bool operator==(const iterator& rhs) const;
+			bool operator!=(const iterator& rhs) const;
+			reference operator*() const;
+
+		private:
+			iterator(MemoryPool* owner, std::size_t blockIndex, std::size_t localIndex);
+
+			std::size_t m_blockIndex;
+			std::size_t m_localIndex;
+			MemoryPool* m_owner;
+	};
 }
 
 #include <Nazara/Core/MemoryPool.inl>

include/Nazara/Core/MemoryPool.inl

@@ -64,7 +64,7 @@ namespace Nz
 			break;
 		}
 
-		if (blockIndex == m_blocks.size())
+		if (blockIndex >= m_blocks.size())
 		{
 			// No more room, allocate a new block
 			blockIndex = m_blocks.size();
@@ -77,6 +77,7 @@ namespace Nz
 
 		auto& block = m_blocks[blockIndex];
 		block.freeEntries.Reset(localIndex);
+		block.occupiedEntries.Set(localIndex);
 		block.occupiedEntryCount++;
 
 		T* entry = reinterpret_cast<T*>(&block.memory[localIndex]);
@@ -117,17 +118,15 @@ namespace Nz
 		std::size_t blockIndex = index / m_blockSize;
 		std::size_t localIndex = index % m_blockSize;
 
-		assert(blockIndex < m_blocks.size());
+		T* entry = GetAllocatedPointer(blockIndex, localIndex);
-		auto& block = m_blocks[blockIndex];
+		PlacementDestroy(entry);
-		assert(!block.freeEntries.Test(localIndex));
 
+		auto& block = m_blocks[blockIndex];
 		assert(block.occupiedEntryCount > 0);
 		block.occupiedEntryCount--;
 
-		T* entry = reinterpret_cast<T*>(&block.memory[localIndex]);
-		PlacementDestroy(entry);
-
 		block.freeEntries.Set(localIndex);
+		block.occupiedEntries.Reset(localIndex);
 	}
 
 	/*!
@@ -192,20 +191,36 @@ namespace Nz
 			if (block.occupiedEntryCount == 0)
 				continue;
 
-			for (std::size_t localIndex = 0; localIndex < m_blockSize; ++localIndex)
+			for (std::size_t localIndex = block.occupiedEntries.FindFirst(); localIndex != block.occupiedEntries.npos; localIndex = block.occupiedEntries.FindNext(localIndex))
 			{
-				if (!block.freeEntries.Test(localIndex))
+				T* entry = reinterpret_cast<T*>(&m_blocks[blockIndex].memory[localIndex]);
-				{
+				PlacementDestroy(entry);
-					T* entry = reinterpret_cast<T*>(&m_blocks[blockIndex].memory[localIndex]);
-					PlacementDestroy(entry);
-				}
 			}
 
-			block.freeEntries.Reset();
+			block.freeEntries.Reset(true);
+			block.occupiedEntries.Reset(false);
 			block.occupiedEntryCount = 0;
 		}
 	}
 
+	/*!
+	* \brief Retrieve an allocated pointer based on a valid entry index
+	*
+	* \param index Entry index
+	*
+	* \return Pointer to the allocated entry
+	*
+	* \remark index must be valid
+	*/
+	template<typename T, std::size_t Alignment>
+	T* MemoryPool<T, Alignment>::RetrieveFromIndex(std::size_t index)
+	{
+		std::size_t blockIndex = index / m_blockSize;
+		std::size_t localIndex = index % m_blockSize;
+
+		return GetAllocatedPointer(blockIndex, localIndex);
+	}
+
 	/*!
 	* \brief Retrieve an entry index based on an allocated pointer
 	*
@@ -240,13 +255,130 @@ namespace Nz
 		return blockIndex * m_blockSize + localIndex;
 	}
 
+	template<typename T, std::size_t Alignment>
+	auto MemoryPool<T, Alignment>::begin() -> iterator
+	{
+		auto [blockIndex, localIndex] = GetFirstAllocatedEntry();
+		return iterator(this, blockIndex, localIndex);
+	}
+
+	template<typename T, std::size_t Alignment>
+	auto MemoryPool<T, Alignment>::end() -> iterator
+	{
+		return iterator(this, InvalidIndex, InvalidIndex);
+	}
+
+	template<typename T, std::size_t Alignment>
+	std::size_t MemoryPool<T, Alignment>::size()
+	{
+		return GetAllocatedEntryCount();
+	}
+
 	template<typename T, std::size_t Alignment>
 	void MemoryPool<T, Alignment>::AllocateBlock()
 	{
 		auto& block = m_blocks.emplace_back();
 		block.freeEntries.Resize(m_blockSize, true);
+		block.occupiedEntries.Resize(m_blockSize, false);
 		block.memory = std::make_unique<AlignedStorage[]>(m_blockSize);
 	}
+
+	template<typename T, std::size_t Alignment>
+	T* MemoryPool<T, Alignment>::GetAllocatedPointer(std::size_t blockIndex, std::size_t localIndex)
+	{
+		assert(blockIndex < m_blocks.size());
+		auto& block = m_blocks[blockIndex];
+		assert(block.occupiedEntries.Test(localIndex));
+
+		return reinterpret_cast<T*>(&block.memory[localIndex]);
+	}
+
+	template<typename T, std::size_t Alignment>
+	std::pair<std::size_t, std::size_t> MemoryPool<T, Alignment>::GetFirstAllocatedEntry() const
+	{
+		return GetFirstAllocatedEntryFromBlock(0);
+	}
+
+	template<typename T, std::size_t Alignment>
+	std::pair<std::size_t, std::size_t> MemoryPool<T, Alignment>::GetFirstAllocatedEntryFromBlock(std::size_t blockIndex) const
+	{
+		// Search in next block
+		std::size_t localIndex = InvalidIndex;
+		for (; blockIndex < m_blocks.size(); ++blockIndex)
+		{
+			auto& block = m_blocks[blockIndex];
+			if (block.occupiedEntryCount == 0)
+				continue;
+
+			localIndex = block.occupiedEntries.FindFirst();
+			assert(localIndex != block.occupiedEntries.npos);
+			break;
+		}
+
+		if (blockIndex >= m_blocks.size())
+			return { InvalidIndex, InvalidIndex };
+
+		return { blockIndex, localIndex };
+	}
+
+	template<typename T, std::size_t Alignment>
+	std::pair<std::size_t, std::size_t> MemoryPool<T, Alignment>::GetNextAllocatedEntry(std::size_t blockIndex, std::size_t localIndex) const
+	{
+		assert(blockIndex < m_blocks.size());
+		auto& block = m_blocks[blockIndex];
+		std::size_t nextLocalIndex = block.occupiedEntries.FindNext(localIndex);
+		if (nextLocalIndex != block.occupiedEntries.npos)
+			return { blockIndex, nextLocalIndex };
+
+		// Search in next block
+		return GetFirstAllocatedEntryFromBlock(blockIndex + 1);
+	}
+
+
+	template<typename T, std::size_t Alignment>
+	MemoryPool<T, Alignment>::iterator::iterator(MemoryPool* owner, std::size_t blockIndex, std::size_t localIndex) :
+	m_blockIndex(blockIndex),
+	m_localIndex(localIndex),
+	m_owner(owner)
+	{
+	}
+
+	template<typename T, std::size_t Alignment>
+	auto MemoryPool<T, Alignment>::iterator::operator++(int) -> iterator
+	{
+		iterator copy(*this);
+		operator++();
+		return copy;
+	}
+
+	template<typename T, std::size_t Alignment>
+	auto MemoryPool<T, Alignment>::iterator::operator++() -> iterator&
+	{
+		auto [blockIndex, localIndex] = m_owner->GetNextAllocatedEntry(m_blockIndex, m_localIndex);
+		m_blockIndex = blockIndex;
+		m_localIndex = localIndex;
+
+		return *this;
+	}
+
+	template<typename T, std::size_t Alignment>
+	bool MemoryPool<T, Alignment>::iterator::operator==(const iterator& rhs) const
+	{
+		assert(m_owner == rhs.m_owner);
+		return m_blockIndex == rhs.m_blockIndex && m_localIndex == rhs.m_localIndex;
+	}
+
+	template<typename T, std::size_t Alignment>
+	bool MemoryPool<T, Alignment>::iterator::operator!=(const iterator& rhs) const
+	{
+		return !operator==(rhs);
+	}
+
+	template<typename T, std::size_t Alignment>
+	auto MemoryPool<T, Alignment>::iterator::operator*() const -> reference
+	{
+		return *m_owner->GetAllocatedPointer(m_blockIndex, m_localIndex);
+	}
 }
 
 #include <Nazara/Core/DebugOff.hpp>

tests/Engine/Core/MemoryPoolTest.cpp

@@ -97,6 +97,23 @@ SCENARIO("MemoryPool", "[CORE][MEMORYPOOL]")
 				CHECK(*vector1 == Nz::Vector2<int>(3, 2));
 				CHECK(*vector2 == Nz::Vector2<int>(3, 5));
 				CHECK(vector3->GetSquaredLength() == Approx(61.f));
+
+				AND_THEN("We iterate on the memory pool")
+				{
+					std::size_t count = 0;
+					int sumX = 0;
+					int sumY = 0;
+					for (T& vec : memoryPool)
+					{
+						count++;
+						sumX += vec.x;
+						sumY += vec.y;
+					}
+
+					CHECK(count == 3);
+					CHECK(sumX == 11);
+					CHECK(sumY == 13);
+				}
 			}
 
 			memoryPool.Reset();
@@ -105,6 +122,12 @@ SCENARIO("MemoryPool", "[CORE][MEMORYPOOL]")
 			CHECK(memoryPool.GetBlockCount() == 2);
 			CHECK(memoryPool.GetFreeEntryCount() == 4);
 
+			bool failure = false;
+			for (T& vec : memoryPool)
+				failure = true;
+
+			CHECK_FALSE(failure);
+
 			memoryPool.Clear();
 			CHECK(allocationCount == 0);
 			CHECK(memoryPool.GetAllocatedEntryCount() == 0);