/* 
 * Copyright (C) 2026 Peter Mikkelsen <petermikkelsen10@gmail.com>
 *
 * This file is part of aplwc.
 *
 * Aplwc is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Aplwc is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with aplwc.  If not, see <https://www.gnu.org/licenses/>.
 */

#include "aplwc_internal.h"

#define ALIGNMENT 8
#define HEADER_SIZE (sizeof(struct chunk) + (sizeof(struct chunk) % ALIGNMENT))

enum chunk_tag {
	CHUNK_TAG_FREE,
	CHUNK_TAG_APLWC_LINE,
	CHUNK_TAG_STRING,
	CHUNK_TAG_TOKEN_LIST,
	CHUNK_TAG_APLWC_VALUE,
	CHUNK_TAG_APLWC_MODULE
};

struct chunk {
	struct chunk *stack;
	size_t size;
	enum chunk_tag tag;
	unsigned int reachable:1;
	unsigned int last:1;
};

void
memory_init(struct aplwc *aplwc)
{
	aplwc->data = NULL;
	aplwc->unaligned_data = NULL;
	aplwc->size = 0;
	aplwc->memptrs = NULL;
	aplwc->refs = NULL;
	aplwc->reachable_bit = 0;
}

void
memory_destroy(struct aplwc *aplwc)
{
	if(aplwc->fns.dealloc && aplwc->unaligned_data)
		aplwc->fns.dealloc(aplwc->unaligned_data);
}

void *
data(struct chunk *c)
{
	uint8_t *p;

	p = (uint8_t*)c;
	return p + HEADER_SIZE;
}

struct chunk *
next(struct chunk *c)
{
	uint8_t *p;

	if(c->last)
		return NULL;

	p = (uint8_t*)c;

	return (struct chunk *)(p + c->size);
}

#define VISIT_FIELD(X, tag, ptr) X(((struct tag *)data(c))->ptr)
#define VISIT_CHUNK_PTRS(X, c) switch(c->tag){ \
	case CHUNK_TAG_APLWC_LINE: \
		VISIT_FIELD(X, aplwc_line, text); \
		VISIT_FIELD(X, aplwc_line, tokens); \
		break; \
	case CHUNK_TAG_STRING: \
		break; \
	case CHUNK_TAG_TOKEN_LIST: \
		VISIT_FIELD(X, token_list, text); \
		break; \
	case CHUNK_TAG_APLWC_VALUE: \
		break; \
	case CHUNK_TAG_APLWC_MODULE: \
		VISIT_FIELD(X, aplwc_module, name); \
		VISIT_FIELD(X, aplwc_module, tokens); \
		VISIT_FIELD(X, aplwc_module, link); \
		break; \
	case CHUNK_TAG_FREE: \
		break; \
	}
#define VISIT_ROOT_PTRS(X) do{ \
	X(aplwc->modules); \
	for(frame = aplwc->memptrs; frame; frame = frame->link) \
		for(i = 0; i < frame->count; i++) \
			if(frame->ptrs[i]) \
				X(*frame->ptrs[i]); \
	for(ref = aplwc->refs; ref; ref = ref->link) \
		if(ref->ptr) \
			X(ref->ptr); \
}while(0)

#define VISIT_ALL_PTRS(X, data) do{ \
	VISIT_ROOT_PTRS(X); \
	for(c = (struct chunk *)(data); c; c = next(c)) \
		VISIT_CHUNK_PTRS(X, c); \
}while(0)

static void
fixrefs(struct aplwc *aplwc, uint8_t *new, uint8_t *old)
{

	struct chunk *c;
	struct memory_ptr_frame *frame;
	struct aplwc_ref *ref;
	size_t add, sub, i;

	if(new > old){
		add = new-old;
		sub = 0;
	}else{
		add = 0;
		sub = old-new;
	}

#define FIX(ptr) do{ \
	if(ptr){ \
		(ptr) = (void*)(((uint8_t*)(ptr))+add); \
		(ptr) = (void*)(((uint8_t*)(ptr))-sub); \
	} \
}while(0)
	VISIT_ALL_PTRS(FIX, new);
#undef FIX
}

static enum aplwc_rc
grow(struct aplwc *aplwc, size_t extra)
{
	uint8_t *unaligned_data, *aligned_data;
	size_t unaligned_newsize, aligned_newsize, adjust;
	struct chunk *free, *head, *last;

	if(aplwc->fns.alloc == NULL)
		goto error;

	unaligned_data = aplwc->fns.alloc(aplwc->size + extra + HEADER_SIZE + 2*ALIGNMENT, &unaligned_newsize);
	if(unaligned_data == NULL)
		goto error;

	adjust = ((uintptr_t)unaligned_data) % ALIGNMENT;
	if(adjust)
		adjust = ALIGNMENT - adjust;
	aligned_newsize = unaligned_newsize - adjust;
	aligned_data = unaligned_data + adjust;

	aligned_newsize -= aligned_newsize % ALIGNMENT;

	assert(aligned_newsize >= (aplwc->size + extra));
	head = (struct chunk *)aligned_data;
	last = NULL;

	if(aplwc->data){
		memmove(aligned_data, aplwc->data, aplwc->size);
		fixrefs(aplwc, aligned_data, aplwc->data);
		memory_destroy(aplwc);
		for(free = head; free; free = next(free))
			if(next(free))
				assert((free)->tag != CHUNK_TAG_FREE);
			else
				last = free;
	}

	free = (struct chunk *)(aligned_data + aplwc->size);
	free->tag = CHUNK_TAG_FREE;
	free->size = aligned_newsize - aplwc->size;
	free->stack = NULL;
	free->last = 1;

	if(last){
		if(last && (last->tag == CHUNK_TAG_FREE))
			last->size += free->size;
		else
			last->last = 0;
	}

	aplwc->unaligned_data = unaligned_data;
	aplwc->data = aligned_data;
	aplwc->size = aligned_newsize;
	return 1;

error:
	aplwc->last_error = APLWC_OUT_OF_MEMORY;
	return 0;
}

static struct chunk *
find_free(struct aplwc *aplwc, size_t min)
{
	struct chunk *c;

	min += HEADER_SIZE;

	if(aplwc->data == NULL)
		return NULL;
	
	for(c = (struct chunk *)aplwc->data; c; c = next(c)){
		if((c->tag == CHUNK_TAG_FREE) && (c->size >= min))
			break;
	}

	return c;
}

static void
push(struct aplwc *aplwc, struct chunk **stack, void *p)
{
	struct chunk *c;

	if(p == NULL)
		return;
	c = (struct chunk *)(((uint8_t*)p)-HEADER_SIZE);
	assert(c->tag != CHUNK_TAG_FREE);

	if(c->reachable == aplwc->reachable_bit)
		return;
	
	assert(c->stack == NULL);
	c->reachable = aplwc->reachable_bit;
	c->stack = *stack;
	*stack = c;
}

static struct chunk *
pop(struct chunk **stack)
{
	struct chunk *c;

	assert(stack);
	assert(*stack);

	c = *stack;
	*stack = c->stack;
	c->stack = NULL;

	return c;
}

static void
collect(struct aplwc *aplwc)
{
	struct chunk *c, *prev;
	struct chunk *stack;
	struct memory_ptr_frame *frame;
	struct aplwc_ref *ref;
	size_t i;

	aplwc->reachable_bit = ~aplwc->reachable_bit;

	/* Mark reachable chunks */

	stack = NULL;

#define PUSH(ptr) push(aplwc, &stack, ptr)
	VISIT_ROOT_PTRS(PUSH);
	while(stack){
		c = pop(&stack);
		VISIT_CHUNK_PTRS(PUSH, c);
#undef PUSH
	}

	/* Free chunks not marked */
	prev = NULL;
	for(c = (struct chunk *)aplwc->data; c; prev = c, c = next(c)){
		if((c->reachable != aplwc->reachable_bit) || (c->tag == CHUNK_TAG_FREE)){
			if(prev && (prev->tag == CHUNK_TAG_FREE)){
				prev->size += c->size;
				prev->last = c->last;
				c = prev;
			}else{
				c->tag = CHUNK_TAG_FREE;
			}
		}
	}
}

static void
swapptr(struct aplwc *aplwc, void *old, void *new)
{
	struct chunk *c;
	struct memory_ptr_frame *frame;
	struct aplwc_ref *ref;
	size_t i;

#define SWAP(ptr) do{ \
	if((ptr) == old) \
		(ptr) = new; \
}while(0)
	VISIT_ALL_PTRS(SWAP, aplwc->data);
#undef SWAP
}

static void
compact(struct aplwc *aplwc)
{
	struct chunk *c, *prev;
	void *oldptr, *newptr;
	size_t size;
	int last;

	for(c = (struct chunk *)aplwc->data, prev = NULL; c; prev = c, c = next(c)){
		if((prev == NULL) || (prev->tag != CHUNK_TAG_FREE))
			continue;

		if(c->tag == CHUNK_TAG_FREE){
			prev->size += c->size;
			prev->last = c->last;
		}else{
			/* swap prev and c */
			oldptr = data(c);
			newptr = data(prev);

			size = prev->size;
			last = c->last;
			
			memmove(prev, c, c->size);

			c = (struct chunk *)(((uint8_t *)prev) + prev->size);
			c->stack = NULL;
			c->tag = CHUNK_TAG_FREE;
			c->size = size;
			c->last = last;
			prev->last = 0;

			swapptr(aplwc, oldptr, newptr);
		}
	}
}

static void
debug_check_mem(struct aplwc *aplwc)
{
	uint64_t size;
	struct chunk *c, *prev;

	size = 0;
	prev = NULL;
	for(c = (struct chunk *)aplwc->data; c; prev = c, c = next(c)){
		assert(((uint8_t*)c) < (aplwc->data+aplwc->size));
		size += c->size;
		assert((c->last == 0) != (size == aplwc->size));
		if(c->tag == CHUNK_TAG_FREE)
			assert((prev == NULL) || (prev->tag != CHUNK_TAG_FREE));
	}
	
	assert(size == aplwc->size);
}

static void *
alloc(struct aplwc *aplwc, size_t size, enum chunk_tag tag)
{
	struct chunk *free;
	struct chunk *new;
	size_t pad;

	int attempt;

	pad = size % ALIGNMENT;
	if(pad)
		size += ALIGNMENT - pad;
		
	attempt = 0;
retry:
	DEBUG(check_mem(aplwc));
	free = FAILCASE(find_free(aplwc, size), NULL);
	if(!free){
		switch(attempt++){
		case 0:
			collect(aplwc);
			goto retry;
		case 1:
			compact(aplwc);
			goto retry;
		case 2:
			if(grow(aplwc, size))
				goto retry;
			/* fall-through */
		default:
			aplwc->last_error = APLWC_OUT_OF_MEMORY;
			return NULL;
		}
	}

	assert(free->tag == CHUNK_TAG_FREE);

	new = free;

	if(free->size > (size + 2*HEADER_SIZE)){
		free = (struct chunk *)(((uint8_t*)free) + HEADER_SIZE + size);
		free->tag = CHUNK_TAG_FREE;
		free->size = new->size - (HEADER_SIZE + size);
		free->stack = NULL;

		if(new->last){
			new->last = 0;
			free->last = 1;
		}

		new->size = size+HEADER_SIZE;
	}

	new->tag = tag;
	new->reachable = aplwc->reachable_bit;

	DEBUG(check_mem(aplwc));
	return data(new);
}

void
memory_root_new(struct aplwc *aplwc, void *p, struct aplwc_ref *ref)
{
	ref->ptr = p;
	ref->link = aplwc->refs;
	ref->aplwc = aplwc;
	aplwc->refs = ref;
}

void *
memory_root_get(struct aplwc_ref *ref)
{
	return ref->ptr;
}

void
memory_root_free(struct aplwc_ref *ref)
{
	struct aplwc *aplwc;
	struct aplwc_ref **root;

	if(!valid_ref(ref))
		return;

	aplwc = ref->aplwc;

	for(root = &aplwc->refs; (*root); root = &(*root)->link){
		if((*root) == ref){
			*root = (*root)->link;
			break;
		}
	}

	clear_ref(ref);
}

void
memptrs_add(struct aplwc *aplwc, struct memory_ptr_frame *frame, void ***ptrs, size_t count)
{
	size_t i;

	frame->ptrs = ptrs;
	frame->link = aplwc->memptrs;	
	frame->count = count;
	aplwc->memptrs = frame;

	for(i = 0; i < count; i++)
		frame->ptrs[i] = NULL;
}

void
memptrs_remove(struct aplwc *aplwc, struct memory_ptr_frame *frame)
{
	assert(aplwc->memptrs == frame);
	aplwc->memptrs = frame->link;
}

struct aplwc_line *
memory_alloc_aplwc_line(struct aplwc *aplwc)
{
	struct aplwc_line *p;

	p = alloc(aplwc, sizeof(struct aplwc_line), CHUNK_TAG_APLWC_LINE);
	return p;
}

struct string *
memory_alloc_string(struct aplwc *aplwc, size_t bytes)
{
	struct string *p;

	p = alloc(aplwc, sizeof(struct string) + bytes, CHUNK_TAG_STRING);

	return p;
}

struct token_list *
memory_alloc_token_list(struct aplwc *aplwc, size_t tokens)
{
	struct token_list *p;

	p = alloc(aplwc, sizeof(struct token_list) + (tokens * sizeof(struct token)), CHUNK_TAG_TOKEN_LIST);
	return p;
}

struct aplwc_value *
memory_alloc_aplwc_value(struct aplwc *aplwc)
{
	struct aplwc_value *p;

	p = alloc(aplwc, sizeof(struct aplwc_value), CHUNK_TAG_APLWC_VALUE);
	return p;
}

struct aplwc_module *
memory_alloc_aplwc_module(struct aplwc *aplwc)
{
	struct aplwc_module *p;

	p = alloc(aplwc, sizeof(struct aplwc_module), CHUNK_TAG_APLWC_MODULE);
	return p;
}

void
clear_ref(struct aplwc_ref *ref)
{
	ref->aplwc = NULL;
	ref->ptr = NULL;
	ref->link = NULL;
}

int
valid_ref(struct aplwc_ref *ref)
{
	return (ref->aplwc != NULL) && (ref->ptr != NULL);
}