/* libheap/malloc.c - An alternative <libmisc/heap.h> implementation
 *
 * Copyright (C) 2025  Luke T. Shumaker <lukeshu@lukeshu.com>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 */

#include <stdint.h> /* for uintptr_t */
#include <string.h> /* for memset() */

#include <libcr/coroutine.h>
#include <libmisc/assert.h>
#include <libmisc/linkedlist.h>
#include <libmisc/macro.h>

#include <libmisc/_intercept.h>
#include <libheap/inspect.h>

/* Configuration.  ************************************************************/

#include "config.h"

#if __unix__
#ifndef CONFIG_HEAP_SIZE
	#error config.h must define CONFIG_HEAP_SIZE
#endif
#endif

#if CONFIG_HEAP_VALGRIND
	#include <valgrind/valgrind.h>
#endif
#if CONFIG_COROUTINE_VALGRIND
	#include <valgrind/memcheck.h>
#endif

/* Platform-specific initialization.  *****************************************/

struct heap {
	void    *base;
	size_t   size;
};

[[gnu::noinline]] [[gnu::weak]] struct heap heap_get(void) {
#if __unix__
	static char heap[CONFIG_HEAP_SIZE] = {};
	return (struct heap){
		.base = heap,
		.size = sizeof(heap),
	};
#else
	extern char __end__;     /* set by linker; end of pre-allocated memory (start of heap) */
	extern char __HeapLimit; /* set by linker; end of heap */
	return (struct heap){
		.base = &__end__,
		.size = &__HeapLimit - &__end__,
	};
#endif
}

/* Core logic.  ***************************************************************/

/* data structures =============================*/

SLIST_DECLARE(waiter_list);
struct waiter {
	cid_t   cid;
	void   *ptr;
	size_t  align;
	size_t  size;
};
SLIST_DECLARE_NODE(waiter_list, struct waiter);

struct record {
	struct record   *next;
	cid_t            owner;
};

static struct {
	struct record           *base;
	size_t                   size;

	struct waiter_list       waiters;
	bool                     unpausing;
} heap_meta = {};

static void heap_ensure_init(void) {
	if (heap_meta.base)
		return;

	struct heap heap = heap_get();
	heap_meta.base = heap.base;
	heap_meta.size = heap.size;

	*(struct record *)(heap.base) = (struct record){
		.next = (void*)((uintptr_t)heap.base + heap.size),
		.owner = HEAP_CID_UNALLOCATED,
	};
}
#define heap_assert_init() assert(heap_meta.base)

static bool heap_record_is_valid(struct record *rec) {
	return (uintptr_t)heap_meta.base <= (uintptr_t)rec &&
		(uintptr_t)&rec[1] <= (uintptr_t)heap_meta.base+heap_meta.size;
}
static void *heap_record_base(struct record *rec) {
	return &rec[1];
}
static size_t heap_record_size(struct record *rec) {
	return (size_t)((uintptr_t)rec->next - (uintptr_t)&rec[1]);
}

#ifndef NDEBUG
static bool heap_ptr_is_valid(void *ptr) {
	return (uintptr_t)&heap_meta.base[1] <= (uintptr_t)ptr &&
		(uintptr_t)ptr <= (uintptr_t)heap_meta.base+heap_meta.size;
}
#endif

/* algorithm ===================================*/

static size_t heap_record_fits(struct record *rec, struct waiter *waiter, bool allow_overlapping_move) {
	if (rec->owner != HEAP_CID_UNALLOCATED && heap_record_base(rec) != waiter->ptr)
		return 0;
	uintptr_t effective_base = (uintptr_t)heap_record_base(rec);
	size_t effective_size = heap_record_size(rec);
	struct record *next = rec->next;
	if (allow_overlapping_move && heap_record_is_valid(next) && heap_record_base(next) == waiter->ptr) {
		effective_size += sizeof(struct record) + heap_record_size(next);
		next = next->next;
	}
	if (heap_record_is_valid(next) && next->owner == HEAP_CID_UNALLOCATED) {
		effective_size += sizeof(struct record) + heap_record_size(next);
		next = next->next;
	}
	if (effective_base % waiter->align) {
		effective_base = LM_ROUND_UP(effective_base + sizeof(struct record), waiter->align);
		size_t lost = effective_base - (uintptr_t)heap_record_base(rec);
		if (lost > effective_size)
			return 0;
		effective_size -= lost;
	}
	if (effective_size < waiter->size)
		return 0;
	return (effective_base+effective_size) - (uintptr_t)rec;
}

static cid_t heap_getcid(void) {
	cid_t ret = cr_maybe_getcid();
	if (!ret)
		ret = HEAP_CID_KERNEL;
	return ret;
}

static bool heap_alloc_can_alloc(struct waiter *waiter, bool check [[maybe_unused]]) {
#ifndef NDEBUG
	uintptr_t other_base = 0;
	size_t    other_size = 0;
	bool      other_found = !check;
#endif
	for (struct record *rec = heap_meta.base; heap_record_is_valid(rec); rec = rec->next) {
#ifndef NDEBUG
		if (!other_found && rec->owner != waiter->cid) {
			if (!other_size)
				other_base = (uintptr_t)rec;
			other_size = (uintptr_t)rec->next - other_base;
			uintptr_t effective_base = other_base + sizeof(struct record);
			size_t effective_size = other_size - sizeof(struct record);
			if (effective_base % waiter->align) {
				effective_base = LM_ROUND_UP(effective_base + sizeof(struct record), waiter->align);
				effective_size -= effective_base - (other_base + sizeof(struct record));
			}
			if (effective_size >= waiter->size)
				other_found = true;
		}
#endif
		if (heap_record_fits(rec, waiter, true))
			return true;
	}
	assert(other_found);
	return false;
}

#define REDZONE 0

void *__lm_heap_aligned_realloc(void *ptr, size_t align, size_t size) {
	heap_ensure_init();

	cid_t owner = heap_getcid();

	/* libmisc/heap.c forces a minimum alignement of
	 * sizeof(void*), so having a size less than LM_ROUNDUP(...,
	 * sizeof(void*)) won't ever save any space, and will actually
	 * cost space as we'll need to insert another record for
	 * padding.
	 */
	size = LM_ROUND_UP(size, sizeof(void*));

	struct record *old_rec = NULL;
	size_t old_size = 0;
	if (ptr) {
		assert(heap_ptr_is_valid(ptr));

		old_rec = heap_meta.base;
		while (heap_record_is_valid(old_rec) && (uintptr_t)heap_record_base(old_rec) < (uintptr_t)ptr)
			old_rec = old_rec->next;
		assert(heap_record_is_valid(old_rec) && heap_record_base(old_rec) == ptr && old_rec->owner == owner);
		old_size = heap_record_size(old_rec);
		if (old_size >= size)
			return ptr;
	}

	/* wait ======================*/
	struct waiter_list_node waiter = { .val = {
		.cid = owner,
		.ptr = ptr,
		.align = align,
		.size = size,
	}};
	slist_push_to_rear(&heap_meta.waiters, &waiter);
#ifdef NDEBUG
	if (heap_meta.waiters.front != &waiter || !heap_alloc_can_alloc(&waiter.val, false))
#else
	bool can_alloc = heap_alloc_can_alloc(&waiter.val, true);
	if (heap_meta.waiters.front != &waiter || !can_alloc)
#endif
		cr_pause_and_yield();
	assert(heap_meta.waiters.front == &waiter && heap_alloc_can_alloc(&waiter.val, false));
	slist_pop_from_front(&heap_meta.waiters);

	/* allocate ==================*/
	/* decision making */
	enum {
		MODE_SIMPLE,
		MODE_REUSE,
		MODE_OVERLAPPING_MOVE,
	} mode;
	struct record *choice = NULL;
	if (old_rec && heap_record_fits(old_rec, &waiter.val, false)) {
		/* reuse the existing block if we can */
		assert(old_rec->next->owner == HEAP_CID_UNALLOCATED);
		assert((uintptr_t)heap_record_base(old_rec) % align == 0);
		choice = old_rec;
		mode = MODE_REUSE;
	} else {
		/* else, use the smallest possible block without doing an overlapping-move */
		size_t choice_size = 0;
		for (struct record *rec = heap_meta.base; heap_record_is_valid(rec); rec = rec->next) {
			size_t rec_size = heap_record_fits(rec, &waiter.val, false);
			if (rec_size && (!choice || rec_size < choice_size)) {
				choice = rec;
				choice_size = rec_size;
				mode = MODE_SIMPLE;
			}
		}
		if (!choice) {
			/* dang, we're going to have to do an overlapping move */
			assert(old_rec);
			for (struct record *rec = heap_meta.base; heap_record_is_valid(rec); rec = rec->next) {
				if (rec->next == old_rec) {
					choice = rec;
					break;
				}
			}
			assert(choice);
			mode = MODE_OVERLAPPING_MOVE;
		}
	}
	assert(heap_record_is_valid(choice));
	/* align */
	if ((uintptr_t)heap_record_base(choice) % align) {
		assert(mode != MODE_REUSE);
		uintptr_t effective_base = LM_ROUND_UP((uintptr_t)heap_record_base(choice) + sizeof(struct record), align);
		struct record *new = (void*)(effective_base - sizeof(struct record));
		*new = (struct record){
			.next = choice->next,
		};
		choice->next = new;
		choice->owner = HEAP_CID_UNALLOCATED;
		choice = new;
	}
	assert(heap_record_is_valid(choice));
	assert((uintptr_t)heap_record_base(choice) % align == 0);
	/* actuate */
	switch (mode) {
	case MODE_OVERLAPPING_MOVE:
		assert(choice->next == old_rec);
		*choice = *old_rec;
		memmove(heap_record_base(choice), heap_record_base(old_rec), old_size);
		[[fallthrough]];
	case MODE_REUSE:
		/* invariant: `choice` is the final record address that we will use */
		if (heap_record_size(choice) < size) {
			assert(choice->next->owner == HEAP_CID_UNALLOCATED);
			choice->next = choice->next->next;
		}
		[[fallthrough]];
	case MODE_SIMPLE:
	}
	/* invariant: `choice` is the final record address that we will use */
	/* invariant: we won't be gobbling any more records */
	if (heap_record_size(choice) > size + sizeof(struct record)) {
		struct record *new = (void*)((uintptr_t)heap_record_base(choice) + size);
		*new = (struct record){
			.next = choice->next,
			.owner = HEAP_CID_UNALLOCATED,
		};
		choice->next = new;
	}
	void *ret = heap_record_base(choice);
	if (old_rec && old_rec != choice)
		memcpy(ret, heap_record_base(old_rec), old_size);
	memset(ret+old_size, 0, heap_record_size(choice)-old_size);
	choice->owner = owner;
	if (mode == MODE_SIMPLE && ptr)
		__lm_heap_free(ptr);

#if CONFIG_HEAP_VALGRIND
	if (old_rec && choice == old_rec)
		VALGRIND_RESIZEINPLACE_BLOCK(ret, old_size, heap_record_size(choice), REDZONE);
	else
		VALGRIND_MALLOCLIKE_BLOCK(ret, heap_record_size(choice), REDZONE, true);
#endif

	/* return ====================*/
	if (heap_meta.waiters.front &&
	    heap_alloc_can_alloc(&heap_meta.waiters.front->val, false))
		cr_unpause(heap_meta.waiters.front->val.cid);
	else
		heap_meta.unpausing = false;
	return ret;
}

void __lm_heap_free(void *ptr) {
	heap_assert_init();
	assert(heap_ptr_is_valid(ptr));

	struct record *prev2 = NULL;
	struct record *prev = NULL;
	struct record *rec = heap_meta.base;
	while (rec && (uintptr_t)heap_record_base(rec) < (uintptr_t)ptr) {
		prev2 = prev;
		prev = rec;
		rec = rec->next;
	}
	assert(rec && heap_record_base(rec) == ptr);
	/*assert(rec->owner == heap_getcid());*/

#if CONFIG_COROUTINE_VALGRIND
	VALGRIND_MAKE_MEM_UNDEFINED(ptr, heap_record_size(rec));
#endif

	if (prev && prev->owner == HEAP_CID_UNALLOCATED) {
		prev->next = rec->next;
		memset(rec, 0, sizeof(struct record));
		rec = prev;
	} else {
		rec->owner = HEAP_CID_UNALLOCATED;
	}
	if (heap_record_is_valid(rec->next) && rec->next->owner == HEAP_CID_UNALLOCATED) {
		struct record *onext = rec->next;
		rec->next = rec->next->next;
		memset(onext, 0, sizeof(struct record));
	}

#if CONFIG_HEAP_VALGRIND
	VALGRIND_FREELIKE_BLOCK(ptr, REDZONE);
#endif

	if (!heap_meta.unpausing &&
	    heap_meta.waiters.front &&
	    (heap_record_fits(prev2, &heap_meta.waiters.front->val, true) ||
	     heap_record_fits(rec, &heap_meta.waiters.front->val, true))) {
		cr_unpause(heap_meta.waiters.front->val.cid);
		heap_meta.unpausing = true;
	}
}

void __lm_heap_take(void *ptr) {
	if (!ptr)
		return;
	heap_assert_init();
	assert(heap_ptr_is_valid(ptr));

	struct record *rec = heap_meta.base;
	while (heap_record_is_valid(rec) && (uintptr_t)heap_record_base(rec) < (uintptr_t)ptr)
		rec = rec->next;
	assert(heap_record_is_valid(rec) && heap_record_base(rec) == ptr && rec->owner);

	rec->owner = heap_getcid();
}

void heap_usage(struct heap_info *out) {
	assert(out);
	heap_ensure_init();

	memset(out, 0, sizeof(*out));
	for (struct record *rec = heap_meta.base; heap_record_is_valid(rec); rec = rec->next) {
		out->allocated[rec->owner].sum += heap_record_size(rec);
		out->allocated[rec->owner].nseg++;
		out->internal_use += sizeof(struct record);
	}
}