Add chain allocator to the thread local allocator.

2 files changed, 129 insertions, 10 deletions

diff --git a/util/thread_local_allocator.cpp b/util/thread_local_allocator.cpp
index 599ef80..192906c 100644
--- a/util/thread_local_allocator.cpp
+++ b/util/thread_local_allocator.cpp
@@ -17,22 +17,140 @@
  */
 
 #include "thread_local_allocator.hpp"
+#include <assert.h>
+#include <stdint.h>
+#include <memory>
 
 namespace dxil_spv
 {
-unsigned alloc_count;
-void begin_thread_allocator_context()
+static constexpr size_t BLOCK_SIZE = 64 * 1024;
+
+class ChainAllocator
 {
+public:
+	void reset();
+	void *allocate(size_t size);
+
+private:
+	struct MallocDeleter
+	{
+		void operator()(void *ptr)
+		{
+			free(ptr);
+		}
+	};
+
+	struct Block
+	{
+		explicit Block(size_t size);
+		void *allocate(size_t size);
+
+		std::unique_ptr<uint8_t, MallocDeleter> block;
+		size_t offset = 0;
+		size_t block_size = 0;
+	};
+	std::vector<Block> blocks;
+	std::vector<Block> huge_blocks;
+	unsigned block_index = 0;
+
+	bool ensure_block();
+	void *allocate_huge(size_t size);
+};
 
+ChainAllocator::Block::Block(size_t size)
+	: block(static_cast<uint8_t *>(malloc(size))), block_size(size)
+{
 }
 
-void end_thread_allocator_context()
+void *ChainAllocator::Block::allocate(size_t size)
 {
+	offset = (offset + 15) & ~size_t(15);
+	if (offset + size <= block_size)
+	{
+		void *ret = block.get() + offset;
+		offset += size;
+		return ret;
+	}
+	else
+		return nullptr;
+}
 
+static thread_local ChainAllocator *allocator;
+
+void ChainAllocator::reset()
+{
+	for (auto &block : blocks)
+		block.offset = 0;
+	block_index = 0;
+	huge_blocks.clear();
 }
 
-void reset_thread_allocator_context()
+bool ChainAllocator::ensure_block()
+{
+	blocks.emplace_back(BLOCK_SIZE);
+	return bool(blocks.back().block);
+}
+
+void *ChainAllocator::allocate_huge(size_t size)
+{
+	huge_blocks.emplace_back(size);
+	return huge_blocks.back().block.get();
+}
+
+void *ChainAllocator::allocate(size_t size)
 {
+	if (size > BLOCK_SIZE)
+		return allocate_huge(size);
+
+	if (block_index >= blocks.size() && !ensure_block())
+		return nullptr;
+
+	void *ptr = blocks[block_index].allocate(size);
+	if (ptr)
+		return ptr;
+
+	block_index++;
+	if (block_index >= blocks.size() && !ensure_block())
+		return nullptr;
 
+	return blocks[block_index].allocate(size);
+}
+
+void *allocate_in_thread(size_t size)
+{
+	if (!allocator)
+		return malloc(size);
+
+	return allocator->allocate(size);
+}
+
+void free_in_thread(void *ptr)
+{
+	if (!allocator)
+	{
+		free(ptr);
+		return;
+	}
+
+	// Don't bother freeing ...
+}
+
+void begin_thread_allocator_context()
+{
+	assert(!allocator);
+	allocator = new ChainAllocator;
+}
+
+void end_thread_allocator_context()
+{
+	assert(allocator);
+	delete allocator;
+	allocator = nullptr;
+}
+
+void reset_thread_allocator_context()
+{
+	assert(allocator);
+	allocator->reset();
 }
 }
diff --git a/util/thread_local_allocator.hpp b/util/thread_local_allocator.hpp
index 2677f1a..2b33f75 100644
--- a/util/thread_local_allocator.hpp
+++ b/util/thread_local_allocator.hpp
@@ -23,10 +23,13 @@
 #include <unordered_map>
 #include <stdlib.h>
 #include <stdio.h>
+#include <stddef.h>
 
 namespace dxil_spv
 {
-extern unsigned alloc_count;
+void *allocate_in_thread(std::size_t size);
+void free_in_thread(void *ptr);
+
 template <typename T>
 class ThreadLocalAllocator
 {
@@ -37,16 +40,14 @@ public:
 	template <typename U>
 	ThreadLocalAllocator(const ThreadLocalAllocator<U> &) noexcept {}
 
-	value_type *allocate(std::size_t n)
+	value_type *allocate(size_t n)
 	{
-		alloc_count++;
-		fprintf(stderr, "Count = %u\n", alloc_count);
-		return static_cast<T *>(::malloc(sizeof(T) * n));
+		return static_cast<value_type *>(allocate_in_thread(sizeof(T) * n));
 	}
 
 	void deallocate(value_type *p, std::size_t)
 	{
-		::free(p);
+		free_in_thread(p);
 	}
 
 	using is_always_equal = std::true_type;