ExternalAllocator

ExternalAllocator.cpp

#include "ExternalAllocator.h"
#include <stdio.h>
#include <stdlib.h>

#define ROUNDDOWN(a, b)	(((a) / (b)) * (b))
#define ROUNDUP(a, b)	ROUNDDOWN((a) + (b) - 1, b)

static void Assert(bool cond) {if (!cond) abort();}


ExternalAllocator::~ExternalAllocator()
{
	fSizeMap.Clear();
	for (;;) {
		Block *block = fAdrMap.LeftMost();
		if (block == NULL) break;
		fAdrMap.Remove(block);
		delete block;
	}
}

void ExternalAllocator::Register(uint64_t ptr, uint64_t size)
{
	Block *block = new Block(ptr, size);
	fAdrMap.Insert(block);
	fSizeMap.Insert(block);
	fTotalSize += size;
}

bool ExternalAllocator::Alloc(uint64_t &ptr, uint64_t size)
{
	Block *block = fSizeMap.FindClosest(size, false);
	if (block == NULL)
		return false;

	fSizeMap.Remove(block);
	if (block->fSize > size) {
		Block *remainingBlock = new Block(block->fAdr + size, block->fSize - size);
		fAdrMap.Insert(remainingBlock);
		fSizeMap.Insert(remainingBlock);
		block->fSize = size;
	}

	block->fAllocated = true;
	ptr = block->fAdr;
	fAllocSize += block->fSize;
	return true;
}

bool ExternalAllocator::AllocAligned(uint64_t &ptr, uint64_t size, uint64_t align)
{
	Block *block = fSizeMap.FindClosest(size + (align - 1), false);
	if (block == NULL)
		return false;

	uint64_t retPtr = ROUNDUP(block->fAdr, align);
	ptr = retPtr;
	return AllocAt(retPtr, size);
}

bool ExternalAllocator::AllocAt(uint64_t ptr, uint64_t size)
{
	Block *block = fAdrMap.FindClosest(ptr, true);

	if (block == NULL || block->fAllocated || !(ptr >= block->fAdr && ptr + size <= block->fAdr + block->fSize))
		return false;

	if (block->fAdr < ptr) {
		Block *left = block;
		block = new Block(ptr, left->fSize - (ptr - left->fAdr));
		fAdrMap.Insert(block);
		fSizeMap.Insert(block);
		fSizeMap.Remove(left);
		left->fSize = ptr - left->fAdr;
		fSizeMap.Insert(left);
	}
	if (block->fSize > size) {
		Block *right = new Block(ptr + size, block->fSize - size);
		fAdrMap.Insert(right);
		fSizeMap.Insert(right);

		fSizeMap.Remove(block);
		block->fSize = size;
		fSizeMap.Insert(block);
	}

	fSizeMap.Remove(block);
	block->fAllocated = true;
	fAllocSize += block->fSize;
	return true;
}

void ExternalAllocator::Free(uint64_t ptr)
{
	Block *block = fAdrMap.Find(ptr);
	if (block == NULL || !block->fAllocated) abort();
	fSizeMap.Insert(block);
	block->fAllocated = false;
	fAllocSize -= block->fSize;

	Block *prev = fAdrMap.Previous(block);
	if (prev != NULL && !prev->fAllocated) {
		fSizeMap.Remove(prev);
		prev->fSize += block->fSize;

		fAdrMap.Remove(block);
		fSizeMap.Remove(block);
		delete block;

		block = prev;
		fSizeMap.Insert(block);
	}

	Block *next = fAdrMap.Next(block);
	if (next != NULL && !next->fAllocated) {
		fSizeMap.Remove(block);
		block->fSize += next->fSize;

		fAdrMap.Remove(next);
		fSizeMap.Remove(next);
		delete next;

		fSizeMap.Insert(block);
	}
}

void ExternalAllocator::Dump()
{
	printf("\n");
	printf("Blocks:\n");
	for (Block *it = fAdrMap.LeftMost(); it != NULL; it = fAdrMap.Next(it)) {
		printf("  %#" PRIx64 ": %#" PRIx64 ": %s\n", it->fAdr, it->fAdr + it->fSize - 1, it->fAllocated ? "alloc" : "free");
	}
	printf("Free:\n");
	for (Block *it = fSizeMap.LeftMost(); it != NULL; it = fSizeMap.Next(it)) {
		printf("  %#" PRIx64 ": %#" PRIx64 "\n", it->fSize, it->fAdr);
	}
}

ExternalAllocator.h

#pragma once
#include <util/AVLTree.h>


class ExternalAllocator {
private:
	class Block {
	public:
		struct AdrNodeDef {
			typedef uint64 Key;
			typedef Block Value;

			inline AVLTreeNode* GetAVLTreeNode(Value* value) const
			{
				return &value->fAdrNode;
			}

			inline Value* GetValue(AVLTreeNode* node) const
			{
				return (Block*)((char*)node - offsetof(Block, fAdrNode));
			}

			inline int Compare(const Key& a, const Value* b) const
			{
				if (a < b->fAdr) return -1;
				if (a > b->fAdr) return 1;
				return 0;
			}

			inline int Compare(const Value* a, const Value* b) const
			{
				if (a->fAdr < b->fAdr) return -1;
				if (a->fAdr > b->fAdr) return 1;
				return 0;
			}
		};

		struct SizeNodeDef {
			typedef uint64 Key;
			typedef Block Value;

			inline AVLTreeNode* GetAVLTreeNode(Value* value) const
			{
				return &value->fSizeNode;
			}

			inline Value* GetValue(AVLTreeNode* node) const
			{
				return (Block*)((char*)node - offsetof(Block, fSizeNode));
			}

			inline int Compare(const Key& a, const Value* b) const
			{
				if (a < b->fSize) return -1;
				if (a > b->fSize) return 1;
				return 0;
			}

			inline int Compare(const Value* a, const Value* b) const
			{
				if (a->fSize < b->fSize) return -1;
				if (a->fSize > b->fSize) return 1;
				if (a->fAdr < b->fAdr) return -1;
				if (a->fAdr > b->fAdr) return 1;
				return 0;
			}
		};

		uint64 fAdr;
		uint64 fSize;
		AVLTreeNode fAdrNode;
		AVLTreeNode fSizeNode;
		bool fAllocated = false;

		Block(uint64 adr, uint64 size): fAdr(adr), fSize(size) {}
	};

	AVLTree<Block::AdrNodeDef> fAdrMap;
	AVLTree<Block::SizeNodeDef> fSizeMap;
	uint64_t fTotalSize, fAllocSize;

public:
	~ExternalAllocator();

	void Register(uint64_t ptr, uint64_t size);

	[[nodiscard]] bool Alloc(uint64_t &ptr, uint64_t size);
	[[nodiscard]] bool AllocAligned(uint64_t &ptr, uint64_t size, uint64_t align);
	[[nodiscard]] bool AllocAt(uint64_t ptr, uint64_t size);
	void Free(uint64_t ptr);

	inline uint64_t TotalSize() {return fTotalSize;}
	inline uint64_t AllocSize() {return fAllocSize;}
	
	void Dump();
};