/* 
 * 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,
	CHUNK_TAG_MODULE_MEMBER,
	CHUNK_TAG_SYMBOL
};

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

#if APLWC_BUILD_TYPE_DEBUG
static void debug_check_mem(struct aplwc *);
#endif /* APLWC_BUILD_TYPE_DEBUG */

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(ALWAYS(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;
	assert((((uintptr_t)p) % ALIGNMENT) == 0);
	p += c->size;
	assert((((uintptr_t)p) % ALIGNMENT) == 0);

	return (struct chunk *)p;
}

#define VISIT(ptr) (ptr) = fn(ptr, aux)

static void
visit_chunk_ptrs(struct chunk *c, void *(*fn)(void *, void *), void *aux)
{
	union {
		void *vp;
		struct aplwc_line *line;
		struct token_list *tokens;
		struct aplwc_module *module;
		struct module_member *module_member;
		struct symbol *symbol;
	} d;

	assert(c != NULL);

	d.vp = data(c);
	switch(c->tag){
	case CHUNK_TAG_FREE:
		break;
	case CHUNK_TAG_APLWC_LINE:
		VISIT(d.line->text);
		VISIT(d.line->text);
		VISIT(d.line->tokens);
		break;
	case CHUNK_TAG_STRING:
		break;
	case CHUNK_TAG_TOKEN_LIST:
		VISIT(d.tokens->text);
		break;
	case CHUNK_TAG_APLWC_VALUE:
		break;
	case CHUNK_TAG_APLWC_MODULE:
		VISIT(d.module->name);
		VISIT(d.module->tokens);
		VISIT(d.module->imports);
		VISIT(d.module->members);
		VISIT(d.module->link);
		break;
	case CHUNK_TAG_MODULE_MEMBER:
		VISIT(d.module_member->module);
		VISIT(d.module_member->real_member);
		VISIT(d.module_member->link);
		break;
	case CHUNK_TAG_SYMBOL:
		VISIT(d.symbol->text);
		VISIT(d.symbol->link);
		break;
	}
}

static void
visit_root_ptrs(struct aplwc *aplwc, void *(*fn)(void *, void *), void *aux)
{
	struct memory_ptr_frame *frame;
	struct aplwc_ref *ref;
	size_t i;

	VISIT(aplwc->modules);
	VISIT(aplwc->symtab);

	for(frame = aplwc->memptrs; frame; frame = frame->link)
		for(i = 0; i < frame->count; i++)
			if(frame->ptrs[i])
				VISIT(*frame->ptrs[i]);

	for(ref = aplwc->refs; ref; ref = ref->link)
		if(ref->ptr)
			VISIT(ref->ptr);
}

static void
visit_all_ptrs(struct aplwc *aplwc, void *(*fn)(void *, void *), void *aux)
{
	struct chunk *c;

	visit_root_ptrs(aplwc, fn, aux);
	for(c = (struct chunk *)aplwc->data; c; c = next(c))
		visit_chunk_ptrs(c, fn, aux);
}

#undef VISIT

struct fixref_args {
	size_t add;
	size_t sub;
};

static void *
fixref(void *p, void *aux)
{
	struct fixref_args *args = aux;
	uint8_t *r;

	if(p == NULL)
		return p;

	r = p;
	
	r += args->add;
	r -= args->sub;
	return r;
}

static void
fixrefs(struct aplwc *aplwc, uint8_t *old)
{
	uint8_t *new = aplwc->data;
	struct fixref_args args;

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

	visit_all_ptrs(aplwc, fixref, &args);
}

static enum aplwc_rc
grow(struct aplwc *aplwc, size_t extra)
{
	uint8_t *unaligned_data, *aligned_data, *old_data;
	size_t unaligned_newsize, aligned_newsize, adjust, old_size;
	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;

	old_data = aplwc->data;
	old_size = aplwc->size;

	aplwc->data = aligned_data;
	aplwc->size = aligned_newsize;

	if(old_data){
		memmove(aplwc->data, old_data, old_size);
		fixrefs(aplwc, old_data);
		memory_destroy(aplwc);
		for(free = head; free; free = next(free))
			if(next(free))
				assert((free)->tag != CHUNK_TAG_FREE);
			else
				last = free;
	}
	aplwc->unaligned_data = unaligned_data; /* must be after memory_destroy */

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

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

	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;
}

struct push_args {
	struct aplwc *aplwc;
	struct chunk **stack;
};

static void *
push(void *p, void *aux)
{
	struct chunk *c;
	struct push_args *args;

	if(p == NULL)
		return p;

	args = aux;

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

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

	return p;
}

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 push_args args;

	aplwc->reachable_bit = ~aplwc->reachable_bit;

	/* Mark reachable chunks */

	stack = NULL;
	args.aplwc = aplwc;
	args.stack = &stack;

	visit_root_ptrs(aplwc, push, &args);
	while(stack){
		c = pop(&stack);
		visit_chunk_ptrs(c, push, &args);
	}

	/* 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;
			}
		}
	}
}

struct swap_args {
	void *old;
	void *new;
};

static void *
swap(void *p, void *aux)
{
	struct swap_args *args;

	args = aux;
	if(p == args->old)
		return args->new;
	else
		return p;
}

static void
swapptr(struct aplwc *aplwc, void *old, void *new)
{
	struct swap_args args;

	args.old = old;
	args.new = new;

	visit_all_ptrs(aplwc, swap, &args);
}

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;
			c = prev;
		}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
shuffle(struct aplwc *aplwc)
{
	struct chunk *c, *prev, *prev_prev;
	size_t move;

	compact(aplwc);

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

	if(prev == NULL) /* No chunks */
		return;

	if(prev->tag != CHUNK_TAG_FREE){ /* No free space at the end */
		if(grow(aplwc, HEADER_SIZE*8))
			shuffle(aplwc);
		return;
	}

	if(prev_prev == NULL) /* Only free space */
		return;

	assert(prev);
	assert(prev->tag == CHUNK_TAG_FREE);
	move = prev->size;
	prev_prev->last = 1;
	memmove(aplwc->data+move, aplwc->data, aplwc->size-move);

	aplwc->data += move;
	aplwc->size -= move;
	fixrefs(aplwc, aplwc->data-move);

	aplwc->data -= move;
	aplwc->size += move;

	c = (struct chunk *)aplwc->data;
	c->stack = NULL;
	c->tag = CHUNK_TAG_FREE;
	c->size = move;
	c->last = 0;
}

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


	FAILCASE((void)0, shuffle(aplwc));

	pad = size % ALIGNMENT;
	if(pad)
		size += ALIGNMENT - pad;
		
	attempt = 0;
retry:
	DEBUG(check_mem(aplwc));
	new = 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;
			goto end;
		}
	}

	assert(free->tag == CHUNK_TAG_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;
		}else{
			free->last = 0;
		}

		new->size = size+HEADER_SIZE;
	}

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

end:
	DEBUG(check_mem(aplwc));
	return new ? data(new) : NULL;
}

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;
}

struct module_member *
memory_alloc_module_member(struct aplwc *aplwc)
{
	struct module_member *p;

	p = alloc(aplwc, sizeof(struct module_member), CHUNK_TAG_MODULE_MEMBER);
	return p;
}

struct symbol *
memory_alloc_symbol(struct aplwc *aplwc)
{
	struct symbol *p;

	p = alloc(aplwc, sizeof(struct symbol), CHUNK_TAG_SYMBOL);
	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);
}

#if APLWC_BUILD_TYPE_DEBUG
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));
		assert((c->size % ALIGNMENT) == 0);
		assert(c->stack == NULL);
		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);
}
#endif /* APLWC_BUILD_TYPE_DEBUG */