/*
 * Copyright 2025, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *		Augustin Cavalier <waddlesplash>
 */

#include "PagesAllocator.h"

#include <cstdio>
#include <cstring>
#include <errno.h>
#include <new>
#include <sys/mman.h>

#include <locks.h>
#include <syscalls.h>
#include <util/SplayTree.h>


/*! The local page size. Must be a multiple of the system page size. */
static const size_t kPageSize = B_PAGE_SIZE;

/*! The "largest useful" chunk size: any allocations larger than this
 * will get their own area, rather than sharing the common one(s). */
static const size_t kLargestUsefulChunk = 512 * kPageSize;

/*! Amount of virtual address space to reserve when creating new areas. */
/* (The Haiku kernel will ignore reservations if there's no other address
 * space left, so using a large value here should not hurt.) */
static const size_t kReserveAddressSpace = 1 * 1024 * 1024 * 1024;

/*! Cache up to this many percentage points of free memory (compared to used.) */
static const size_t kFreePercentage = 25;

/*! Always forbid more free memory than this, even if it's below kFreePercentage. */
static const size_t kFreeMaximum = 128 * 1024 * 1024;

/*! Always allow this much free memory, even if it's above kFreePercentage. */
static const size_t kFreeMinimum = 128 * kPageSize;

/*! When a free chunk encompasses a whole area and is smaller than this, unmap it. */
static const size_t kMinimumFreeAreaSize = 16 * kPageSize;


namespace {

class PagesAllocator {
	struct FreeChunk;

public:
	PagesAllocator()
	{
		mutex_init(&fLock, "PagesAllocator lock");
		fUsed = fFree = 0;
		fLastArea = -1;
	}

	~PagesAllocator()
	{
	}

	void BeforeFork()
	{
		mutex_lock(&fLock);
	}

	void AfterFork(bool parent)
	{
		if (parent) {
			mutex_unlock(&fLock);
		} else {
			if (fLastArea >= 0)
				fLastArea = area_for((void*)(fLastAreaTop - 1));

			mutex_init(&fLock, "PagesAllocator lock");
		}
	}

	status_t AllocatePages(size_t allocate, void*& address, bool& cleared)
	{
		MutexLocker locker(fLock);

		if (allocate > kLargestUsefulChunk) {
			// Create an area just for this allocation.
			locker.Unlock();
			area_id area = create_area("heap large allocation",
				&address, B_ANY_ADDRESS, allocate,
				B_NO_LOCK, B_READ_AREA | B_WRITE_AREA);
			if (area >= 0) {
				locker.Lock();
				fUsed += allocate;
				cleared = true;
				return B_OK;
			}
			return area;
		}

		if (fFree >= allocate) {
			// Try to use memory from the cache.
			FreeChunk* chunk = fChunksBySizeTree.FindClosest(allocate, true, true);
			if (chunk != NULL) {
				address = _Use(chunk, allocate);
				cleared = false;
				return B_OK;
			}
		}

		// Not enough memory in the cache. Allocate some more.
		FreeChunk* chunk;
		bool mostlyClear = false;
		status_t status = _Map(allocate, chunk, mostlyClear);
		if (status != B_OK)
			return status;

		address = _Use(chunk, allocate);
		cleared = mostlyClear;
		return B_OK;
	}

	status_t AllocatePagesAt(void* _address, size_t allocate, bool& cleared)
	{
		const addr_t address = (addr_t)_address;
		MutexLocker locker(fLock);

		if (allocate <= kLargestUsefulChunk) {
			FreeChunk* chunk = fChunksByAddressTree.FindClosest(address, false, true);
			if (chunk != NULL && chunk->NextAddress() > address) {
				// The address is in a free chunk.
				size_t remainingAfter = chunk->size - (address - (addr_t)chunk);
				if (remainingAfter < allocate) {
					if (chunk->NextAddress() != fLastAreaTop)
						return B_NO_MEMORY;

					size_t add = allocate - remainingAfter;
					status_t status = _ResizeLastArea(add);
					if (status != B_OK)
						return status;

					chunk = _Insert((void*)chunk->NextAddress(), add);
				}

				// Cut the beginning?
				if (address == (addr_t)chunk) {
					_Use(chunk, allocate);
					cleared = false;
					return B_OK;
				}

				// Why would anyone want to allocate such a specific address, anyway?
				debugger("PagesAllocator: middle-chunk allocation not implemented!");
				return B_ERROR;
			}
		}

		if (address == fLastAreaTop) {
			if (allocate > kLargestUsefulChunk)
				return B_NO_MEMORY;

			status_t status = _ResizeLastArea(allocate);
			if (status == B_OK) {
				fUsed += allocate;
				cleared = true;
				return B_OK;
			}
			return status;
		}
		if (address >= (fLastAreaTop - fLastAreaSize) && address < fLastAreaTop) {
			// We didn't find a matching free chunk, so that isn't going to work.
			return EEXIST;
		}

		locker.Unlock();

		// One last try: see if we can resize the area this belongs to.
		area_info info;
		info.area = area_for((void*)(address - 1));
		if (info.area < 0)
			return info.area;

		status_t status = get_area_info(info.area, &info);
		if (status != B_OK)
			return status;

		if (((addr_t)info.address + info.size) != address)
			return B_NO_MEMORY;

		status = resize_area(info.area, info.size + allocate);
		if (status == B_OK) {
			locker.Lock();
			fUsed += allocate;
			cleared = true;
			return B_OK;
		}

		// TODO: We could add a "resize allowing move" feature to resize_area,
		// which would avoid having to memcpy() the data of large allocations.
		return status;
	}

	status_t FreePages(void* _address, size_t size)
	{
		MutexLocker locker(fLock);

		if (size > fUsed)
			debugger("PagesAllocator: request to free more than allocated");
		fUsed -= size;

		FreeChunk* freedChunk = _Insert(_address, size);

		if (size > kLargestUsefulChunk) {
			// TODO: This doesn't deal with a free of a smaller number of pages
			// within a "large allocation" area. We should probably free those
			// immediately also, or at least mark them specially in the trees
			// (so they don't get reused for anything but the large allocation.)

			locker.Detach();
			return _UnlockingRemoveAndUnmap(freedChunk, false);
		}

		if (!_InLastArea(freedChunk) && freedChunk->size < kMinimumFreeAreaSize) {
			// If this chunk covers a whole area, or is at the head or tail of one
			// that's smaller than kMinimumFreeAreaSize, we unmap it immediately.
			// TODO: Find a way to avoid invoking syscalls here, at least in some cases?
			area_info info = {};
			get_area_info(area_for(freedChunk), &info);
			if (info.size == size || (info.size < kMinimumFreeAreaSize
					&& (info.address == freedChunk
						|| ((addr_t)info.address + info.size) == freedChunk->NextAddress()))) {
				locker.Detach();
				return _UnlockingRemoveAndUnmap(freedChunk, false);
			}
		}

		if (fFree <= _FreeLimit()) {
			// TODO: If not decommitting/unmapping, we might free the pages
			// using MADV_FREE (perhaps on medium-sized chunks with others
			// of equal sizes already present in the tree, at least.)
			return B_OK;
		}

		while (fFree > _FreeLimit()) {
			FreeChunk* chunk = fChunksBySizeTree.FindMax();

			// Try to avoid unmapping the just-freed chunk, or any in the last area
			// besides the topmost.
			while (chunk == freedChunk || _InLastArea(chunk)) {
				if (chunk->NextAddress() == fLastAreaTop)
					break;

				FreeChunk* previousChunk = fChunksBySizeTree.PreviousDontSplay(chunk);
				if (previousChunk == NULL) {
					// Just free the maximum-size chunk, whatever it was.
					chunk = fChunksBySizeTree.FindMax();
					break;
				}
				chunk = previousChunk;
			}

			status_t status = _UnlockingRemoveAndUnmap(chunk);
			if (status != B_OK)
				return status;
		}

		return B_OK;
	}

private:
	void* _Use(FreeChunk* chunk, size_t amount)
	{
		fChunksBySizeTree.Remove(chunk);
		fChunksByAddressTree.Remove(chunk);

		if (chunk->size == amount) {
			// The whole chunk will be used.
			// Nothing special to do in this case.
		} else {
			// Some will be left over.
			// Break the remainder off and reinsert into the trees.
			FreeChunk* newChunk = (FreeChunk*)((addr_t)chunk + amount);
			newChunk->size = (chunk->size - amount);

			fChunksBySizeTree.Insert(newChunk);
			fChunksByAddressTree.Insert(newChunk);
		}

		fUsed += amount;
		fFree -= amount;

		memset(chunk, 0, sizeof(FreeChunk));
		return (void*)chunk;
	}

	FreeChunk* _Insert(void* _address, size_t size)
	{
		fFree += size;

		const addr_t address = (addr_t)_address;
		FreeChunk* chunk;

		FreeChunk* preceding = fChunksByAddressTree.FindClosest(address, false, false);
		if (preceding != NULL && preceding->NextAddress() == address) {
			fChunksBySizeTree.Remove(preceding);
			chunk = preceding;
			chunk->size += size;
		} else {
			chunk = (FreeChunk*)_address;
			chunk->size = size;
			fChunksByAddressTree.Insert(chunk);
		}

		FreeChunk* following = chunk->address_tree_list_link;
		if (following != NULL && chunk->NextAddress() == (addr_t)following) {
			fChunksBySizeTree.Remove(following);
			fChunksByAddressTree.Remove(following);
			chunk->size += following->size;
		}
		fChunksBySizeTree.Insert(chunk);

		return chunk;
	}

	size_t _FreeLimit() const
	{
		size_t freeLimit = ((fUsed * kFreePercentage) / 100);
		if (freeLimit < kFreeMinimum)
			freeLimit = kFreeMinimum;
		else if (freeLimit > kFreeMaximum)
			freeLimit = kFreeMaximum;
		else
			freeLimit = (freeLimit + (kPageSize - 1)) & ~(kPageSize - 1);
		return freeLimit;
	}

private:
	status_t _Map(size_t allocate, FreeChunk*& allocated, bool& mostlyClear)
	{
		if (fLastArea >= 0) {
			addr_t oldTop = fLastAreaTop;
			status_t status = _ResizeLastArea(allocate);
			if (status == B_OK) {
				allocated = _Insert((void*)oldTop, allocate);
				mostlyClear = (allocated == (void*)oldTop);
				return B_OK;
			}
		}

		// Create a new area.
		// TODO: We could use an inner lock here to avoid contention.
		addr_t newAreaBase;
		status_t status = _kern_reserve_address_range(&newAreaBase,
			B_RANDOMIZED_ANY_ADDRESS, kReserveAddressSpace);
		size_t newReservation = (status == B_OK) ? kReserveAddressSpace : 0;

		status = create_area("heap area", (void**)&newAreaBase,
			(status == B_OK) ? B_EXACT_ADDRESS : B_RANDOMIZED_ANY_ADDRESS,
			allocate, B_NO_LOCK, B_READ_AREA | B_WRITE_AREA);
		if (status < B_OK)
			return status;

		if (fLastAreaReservedTop > fLastAreaTop)
			_kern_unreserve_address_range(fLastAreaTop, fLastAreaReservedTop - fLastAreaTop);

		fLastArea = status;
		fLastAreaTop = newAreaBase + allocate;
		fLastAreaSize = allocate;
		fLastAreaReservedTop = newAreaBase + newReservation;

		allocated = _Insert((void*)newAreaBase, allocate);
		mostlyClear = (allocated == (void*)newAreaBase);
		return B_OK;
	}

	bool _InLastArea(const FreeChunk* chunk) const
	{
		return ((addr_t)chunk < fLastAreaTop
			&& (addr_t)chunk >= (fLastAreaTop - fLastAreaSize));
	}

	status_t _ResizeLastArea(size_t amount)
	{
		// TODO: We could use an inner lock here to avoid contention.
		status_t status = resize_area(fLastArea, fLastAreaSize + amount);
		if (status == B_OK) {
			fLastAreaTop += amount;
			fLastAreaSize += amount;
		}
		return status;
	}

	status_t _UnlockingRemoveAndUnmap(FreeChunk* chunk, bool relock = true)
	{
		fChunksByAddressTree.Remove(chunk);
		fChunksBySizeTree.Remove(chunk);
		fFree -= chunk->size;

		status_t status = _UnlockingUnmap(chunk);
		if (relock)
			mutex_lock(&fLock);
		if (status != B_OK) {
			if (!relock)
				mutex_lock(&fLock);

			// Unmap failed; reinsert the chunk.
			fChunksByAddressTree.Insert(chunk);
			fChunksBySizeTree.Insert(chunk);
			fFree += chunk->size;

			if (!relock)
				mutex_unlock(&fLock);
		}
		return status;
	}

	status_t _UnlockingUnmap(FreeChunk* chunk)
	{
		// TODO: We could use an inner lock here to avoid contention.
		MutexLocker locker(fLock, true);

		const size_t size = chunk->size;
		const addr_t address = (addr_t)chunk;
		const addr_t top = address + size;

		const addr_t lastAreaBase = (fLastAreaTop - fLastAreaSize);
		if (address <= lastAreaBase && top >= fLastAreaTop) {
			// The whole area is being deleted.
			const addr_t lastTop = fLastAreaTop, reservedTop = fLastAreaReservedTop;
			fLastArea = -1;
			fLastAreaTop = fLastAreaSize = fLastAreaReservedTop = 0;
			locker.Unlock();

			if (reservedTop > lastTop)
				_kern_unreserve_address_range(lastTop, reservedTop - lastTop);
		} else if (top == fLastAreaTop) {
			// Shrink the top.
			status_t status = resize_area(fLastArea, fLastAreaSize - size);
			if (status != B_OK)
				return status;

			fLastAreaSize -= size;
			fLastAreaTop -= size;
			return B_OK;
		} else if (address == lastAreaBase) {
			// Shrink the bottom.
			if (munmap(chunk, size) != 0)
				return errno;
			fLastAreaSize -= size;
			return B_OK;
		} else if (address >= lastAreaBase && address < fLastAreaTop) {
			// Cut the middle and get the new ID.
			if (munmap(chunk, size) != 0)
				return errno;

			fLastAreaSize = fLastAreaTop - top;
			fLastArea = area_for((void*)(fLastAreaTop - 1));
			return B_OK;
		} else {
			// Not in the last area.
			locker.Unlock();
		}

		if (munmap(chunk, size) != 0)
			return errno;
		return B_OK;
	}

private:
	struct FreeChunk {
		SplayTreeLink<FreeChunk> address_tree_link;
		SplayTreeLink<FreeChunk> size_tree_link;
		FreeChunk* address_tree_list_link;
		FreeChunk* size_tree_list_link;

		size_t size;

	public:
		inline addr_t NextAddress() const { return ((addr_t)this + size); }
	};

	struct ChunksByAddressTreeDefinition {
		typedef addr_t		KeyType;
		typedef FreeChunk	NodeType;

		static addr_t GetKey(const FreeChunk* node)
		{
			return (addr_t)node;
		}

		static SplayTreeLink<FreeChunk>* GetLink(FreeChunk* node)
		{
			return &node->address_tree_link;
		}

		static int Compare(const addr_t& key, const FreeChunk* node)
		{
			if (key == (addr_t)node)
				return 0;
			return (key < (addr_t)node) ? -1 : 1;
		}

		static FreeChunk** GetListLink(FreeChunk* node)
		{
			return &node->address_tree_list_link;
		}
	};
	typedef IteratableSplayTree<ChunksByAddressTreeDefinition> ChunksByAddressTree;

	struct ChunksBySizeTreeDefinition {
		struct KeyType {
			size_t size;
			addr_t address;

		public:
			KeyType(size_t _size) : size(_size), address(0) {}
			KeyType(const FreeChunk* chunk) : size(chunk->size), address((addr_t)chunk) {}
		};
		typedef FreeChunk	NodeType;

		static KeyType GetKey(const FreeChunk* node)
		{
			return KeyType(node);
		}

		static SplayTreeLink<FreeChunk>* GetLink(FreeChunk* node)
		{
			return &node->size_tree_link;
		}

		static int Compare(const KeyType& key, const FreeChunk* node)
		{
			if (key.size == node->size)
				return ChunksByAddressTreeDefinition::Compare(key.address, node);
			return (key.size < node->size) ? -1 : 1;
		}

		static FreeChunk** GetListLink(FreeChunk* node)
		{
			return &node->size_tree_list_link;
		}
	};
	typedef IteratableSplayTree<ChunksBySizeTreeDefinition> ChunksBySizeTree;

private:
	mutex	fLock;

	size_t	fUsed;
	size_t	fFree;

	ChunksByAddressTree	fChunksByAddressTree;
	ChunksBySizeTree	fChunksBySizeTree;

	area_id		fLastArea;
	addr_t		fLastAreaTop;
	size_t		fLastAreaSize;
	size_t		fLastAreaReservedTop;
};

} // namespace


static char sPagesAllocatorStorage[sizeof(PagesAllocator)]
#if defined(__GNUC__) && __GNUC__ >= 4
	__attribute__((__aligned__(alignof(PagesAllocator))))
#endif
	;
static PagesAllocator* sPagesAllocator;


void
__init_pages_allocator()
{
	sPagesAllocator = new(sPagesAllocatorStorage) PagesAllocator;
}


void
__pages_allocator_before_fork()
{
	sPagesAllocator->BeforeFork();
}


void
__pages_allocator_after_fork(int parent)
{
	sPagesAllocator->AfterFork(parent);
}


status_t
__allocate_pages(void** address, size_t length, int flags, uint8* cleared)
{
	if ((length % kPageSize) != 0)
		debugger("PagesAllocator: incorrectly sized allocate");

	if ((flags & MAP_FIXED) != 0)
		return sPagesAllocator->AllocatePagesAt(*address, length, *(bool*)cleared);

	//fprintf(stderr, "AllocatePages! 0x%x\n", (int)length);
	return sPagesAllocator->AllocatePages(length, *address, *(bool*)cleared);
}


status_t
__free_pages(void* address, size_t length)
{
	if ((length % kPageSize) != 0)
		debugger("PagesAllocator: incorrectly sized free");

	//fprintf(stderr, "FreePages! %p, 0x%x\n", address, (int)length);
	sPagesAllocator->FreePages(address, length);

	// Errors returned from FreePages() indicate problems unmapping memory;
	// in case of failure the pages will always be added to the free lists.
	// So we never tell callers about failures here.
	return B_OK;
}