/* Copyright (C) 2024 Peter McGoron
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program 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 Lesser General Public
 * License along with this program. If not, see
 * <https://www.gnu.org/licenses/>.
*/
#include <stdio.h>

/* XXX: Currently valgrind marks the header region as valid for all
 * functions, when it should only be valid inside allocator calls.
 */
#ifdef UNS_VALGRIND
# include <valgrind/valgrind.h>
# include <valgrind/memcheck.h>
# define REDZONE 16
#else
# define REDZONE 0
# define VALGRIND_CREATE_MEMPOOL(pool, rzB, is_zeroed) (void)0
# define VALGRIND_DESTROY_MEMPOOL(pool) (void)0
# define VALGRIND_MEMPOOL_ALLOC(pool, ptr, siz) (void)0
# define VALGRIND_CHECK_MEM_IS_ADDRESSABLE(p, len) 0
# define VALGRIND_MAKE_MEM_DEFINED(p, len) (void)0
# define VALGRIND_MAKE_MEM_NOACCESS(p, len) (void)0
#endif

#include <stdlib.h>
#include <limits.h>
#include <assert.h>
#include <string.h>
#include "uns.h"
#include "cheney_c89.h"

struct ctx {
	/** Pointer to the beginning of the heap. */
	unsigned char *tospace;

	/** Pointer to one past the end of the heap. */
	unsigned char *tospace_end;

	/** Pointer to the next place to alloc data. This may be one
	  * past the end of the heap, meaning there is no space left
	  * on the heap.
	  */
	unsigned char *tospace_alloc;

	/** A value set by the user to control the next heap size after
	  * a collection.
	  */
	size_t new_heap_size;

	uns_cheney_c89_collect_callback cb;
	struct uns_cheney_c89_statistics stats;
};

void uns_cheney_c89_deinit(Uns_GC gc)
{
	struct ctx *ctx = uns_ctx(gc);

	free(ctx->tospace);
	free(ctx);
}

void uns_cheney_c89_set_collect_callback(
	Uns_GC gc,
	uns_cheney_c89_collect_callback cb
)
{
	struct ctx *ctx = uns_ctx(gc);
	ctx->cb = cb;
}

void uns_cheney_c89_set_new_heap_size(Uns_GC gc, size_t l)
{
	struct ctx *ctx = uns_ctx(gc);
	ctx->new_heap_size = l;
}

size_t uns_cheney_c89_get_new_heap_size(Uns_GC gc)
{
	struct ctx *ctx = uns_ctx(gc);
	return ctx->new_heap_size;
}

/** Header of a allocated region is a single uns_sword.
  *
  * 0: Relocated. What follows is a pointer to the relocated region.
  * Positive: allocated bytes.
  * Negative: allocated record of pointers with this many bytes.
  *
  * All lengths must be representable in positive and negative.
  * Hence UNS_SWORD_MIN and UNS_SWORD_MAX are disallowed length
  * values.
  *
  * Relocated pointers only exist during copying and are not present
  * during normal execution.
  */

/** Destructured header. "len" is always the readable length of the
  * region in bytes (it is always positive).
  */
struct hdr {
	enum {
		RELO,
		REGION,
		RECORD
	} typ;
	uns_sword len;
};

/** Number of fields in a record. */
#define REC_FIELDS(n) ((n)/sizeof(void*))

/** Size in bytes of the header. */
#define HDR_LEN sizeof(uns_sword)

/** Extract a header from a pointer to a region. */
#define HDR_PTR(p) ((unsigned char *)p - HDR_LEN)

/** Minimum size of a region. */
#define MIN_REG_LEN sizeof(void*)

/** Destructure header pointed to by p
  *
  * # Parameters
  * - `out notnull hdr`: Destructured header.
  * - `notnull p`: Pointer to a header. This is not a pointer to a region.
  *   For a pointer to a region, use `hdr_extract`.
  */
static void hdr_read(struct hdr *hdr, unsigned char *p)
{
	assert(hdr);
	assert(p);

	if (VALGRIND_CHECK_MEM_IS_ADDRESSABLE(p + HDR_LEN, 1) != 0)
		abort();

	VALGRIND_MAKE_MEM_DEFINED(p, HDR_LEN);
	memcpy(&hdr->len, p, HDR_LEN);
	VALGRIND_MAKE_MEM_NOACCESS(p, HDR_LEN);

	if (hdr->len < 0) {
		hdr->typ = RECORD;
		hdr->len = -hdr->len;
	} else if (hdr->len == 0) {
		hdr->typ = RELO;
		hdr->len = sizeof(void*);
	} else {
		hdr->typ = REGION;
	}

	if (VALGRIND_CHECK_MEM_IS_ADDRESSABLE(p + HDR_LEN, hdr->len) != 0)
		abort();
}

/** Destructure header from a pointer to a region.
  *
  * # Parameters
  * - `out notnull hdr`: Destructured header.
  * - `notnull p`: Pointer to a region. This is not a pointer to the
  *   header.
  */
#define hdr_extract(h,p) hdr_read(h, HDR_PTR(p))

/** Write a header to a location.
  *
  * # Parameters
  * - `notnull hdr`: Header description with all fields filled out. This
  *   function does not do sanity checking and may overflow if given
  *   bad data.
  * - `notnull p`: Header to where the header should be written to.
  *   This will overwrite data at the pointer location. The pointer
  *   becomes a pointer to a header, not a pointer to a region.
  */
static void hdr_write_direct(struct hdr *hdr, unsigned char *p)
{
	uns_sword s;

	assert(hdr);
	assert(p);

	switch (hdr->typ) {
	case REGION: s = hdr->len; break;
	case RELO: s = 0; break;
	case RECORD: s = -hdr->len; break;
	}

	if (VALGRIND_CHECK_MEM_IS_ADDRESSABLE(p + HDR_LEN, hdr->len) != 0)
		abort();

	VALGRIND_MAKE_MEM_DEFINED(p, HDR_LEN);
	memcpy(p, &s, HDR_LEN);
	VALGRIND_MAKE_MEM_NOACCESS(p, HDR_LEN);
}

/** Write to the header of a region.
  *
  * # Parameters
  * - `notnull hdr`: See `hdr_write_direct`.
  * - `notnull p`: Header to a region. The call will do pointer arithmetic
  *   to get the pointer to the header, and write to the header. It will
  *   not overwrite region data.
  */
#define hdr_write(h, p) hdr_write_direct(h, HDR_PTR(p))

/** Returns true if there is enough space in the heap for a region with
  * length `bytes`.
  */
static int enough_space(struct ctx *ctx, uns_sword bytes)
{
	return ctx->tospace_end - ctx->tospace_alloc >= bytes + HDR_LEN + REDZONE;
}

/** Allocate region without bounds checking.
  *
  * # Parameters
  * - `len`: Length in bytes of the entire record. This includes the header
  *   region. The length should have been adjusted by the caller to include
  *   the minimum region length.
  * # Returns
  * A pointer to a region.
  */
static void *raw_alloc(struct ctx *ctx, uns_sword len, int is_record)
{
	unsigned char *p;
	struct hdr hdr = {0};

	assert(len >= HDR_LEN + MIN_REG_LEN);

	hdr.len = len - HDR_LEN;
	if (is_record) {
		hdr.typ = RECORD;
	} else {
		hdr.typ = REGION;
	}

	assert(enough_space(ctx, len));

	p = ctx->tospace_alloc;
	VALGRIND_MEMPOOL_ALLOC(ctx->tospace, p + HDR_LEN, len - HDR_LEN);

	hdr_write_direct(&hdr, p);
	ctx->tospace_alloc += len + REDZONE;

	return p + HDR_LEN;
}

/** Move an entire object to the new heap during collection.
  *
  * This function does nothing if `p` is NULL. Otherwise, `p` is a pointer
  * to a region.
  * If `p` was relocated (typ == `RELO`), then this function returns a
  * pointer to the relocated header in tospace.
  * Otherwise, it allocates memory in the tospace, copies the entire region
  * with its header to the tospace, and modifies the region in the fromspace
  * to be a region of type `RELO`.
  */
static unsigned char *relocate(struct ctx *ctx, unsigned char *p)
{
	void *res;
	struct hdr hdr = {0};

	if (!p)
		return NULL;
	hdr_extract(&hdr, p);

	if (hdr.typ == RELO) {
		memcpy(&res, p, sizeof(void*));
		return res;
	}

	assert(hdr.len >= MIN_REG_LEN);

	/* Write entire region to memory */
	res = raw_alloc(ctx, HDR_LEN + hdr.len, hdr.typ == RECORD);
	memcpy(res, p, hdr.len);
	hdr_write(&hdr, res);

	/* Change old pointer to relocation pointer */
	hdr.typ = RELO;
	hdr_write(&hdr, p);
	memcpy(p, &res, sizeof(void*));
	return res;
}

/** Calculate the starting byte index of an element in a record.
  *
  * # Parameters
  * - `notnull p`: Pointer to a region.
  * - `loc`: Index into the region.
  */
static size_t record_index(Uns_ptr p, size_t loc)
{
	struct hdr hdr = {0};

	assert(p);
	hdr_extract(&hdr, p);
	assert(hdr.typ == RECORD);

	/* Turn hdr.len into the number of records in the region */
	assert(loc < hdr.len / sizeof(void*));
	assert(loc < SIZE_MAX/sizeof(void*));
	return loc * sizeof(void*);
}

void *uns_cheney_c89_get(Uns_ptr p, size_t loc,
                         enum uns_fld_type *typ)
{
	void *r;

	if (typ)
		*typ = UNS_POINTER;
	loc = record_index(p, loc);
	memcpy(&r, (unsigned char *)p + loc, sizeof(void*));
	return r;
}

void uns_cheney_c89_set(Uns_ptr p, size_t loc,
                        enum uns_fld_type typ, void *newp)
{
	assert(typ == UNS_POINTER);
	loc = record_index(p, loc);
	memcpy((unsigned char  *)p + loc, &newp, sizeof(void*));
}

/** Relocate each pointer in a record and record its new pointer.
  *
  * # Parameters
  * - `notnull p`: Pointer to a record.
  * - `len`: Number of elements in the record.
  */
static void scan_record(struct ctx *ctx, void *p, size_t len)
{
	size_t i;
	void *newp;

	for (i = 0; i < len; i++) {
		newp = relocate(ctx,
		                uns_cheney_c89_get(p, i, NULL));
		uns_cheney_c89_set(p, i, UNS_POINTER, newp);
	}
}

/** Main section of the copying algorithm. */
static void relocate_everything(Uns_GC gc)
{
	unsigned char *scanptr;
	struct uns_ctr *root;
	struct ctx *ctx = uns_ctx(gc);
	struct hdr hdr = {0};

	/* Relocate roots */
	for (root = uns_roots(gc); root; root = root->next)
		root->p = relocate(ctx, root->p);

	/* Scan the heap until the end of allocated space. If there
	 * is a record at this location, read the record and relocate
	 * all values, and then change each field to the relocated
	 * record pointer.
	 */
	scanptr = ctx->tospace;
	while (scanptr != ctx->tospace_alloc) {
		/* scanptr currently points to the header data. */
		hdr_read(&hdr, scanptr);
		scanptr += HDR_LEN;

		if (hdr.typ == RECORD)
			scan_record(ctx, scanptr, (size_t)hdr.len/sizeof(void*));
		scanptr += hdr.len + REDZONE;
	}

}

int uns_cheney_c89_collect(Uns_GC gc)
{
	/* Save fromspace */
	struct ctx *ctx = uns_ctx(gc);
	unsigned char *fromspace = ctx->tospace;
	unsigned char *fromspace_lim = ctx->tospace_alloc;

	size_t newlen = ctx->new_heap_size;

	ctx->stats.usage_before = fromspace_lim - fromspace;

	/* Bail out immediately if allocation fails. This preserves
	 * the objects as they were.
	 */
	assert(newlen >= fromspace_lim - fromspace);
	ctx->tospace = malloc(newlen);
	VALGRIND_CREATE_MEMPOOL(ctx->tospace, REDZONE, 0);
	if (!ctx->tospace) {
		ctx->tospace = fromspace;
		return 1;
	}

	/* Setup context to be valid for the allocator */
	ctx->tospace_end = ctx->tospace + newlen;
	ctx->tospace_alloc = ctx->tospace;

	relocate_everything(gc);

	VALGRIND_DESTROY_MEMPOOL(fromspace);
	free(fromspace);

	ctx->stats.usage_after = ctx->tospace_alloc - ctx->tospace;
	ctx->stats.collection_number += 1;
	if (ctx->cb)
		ctx->cb(gc, &ctx->stats);

	return 1;
}

static void *alloc(Uns_GC gc, size_t bytes, int is_record)
{
	struct ctx *ctx = uns_ctx(gc);
	uns_sword bytes_as_sword;

	if (bytes >= UNS_SWORD_MAX) {
		uns_on_oom(gc);
		return NULL;
	} else if (bytes < MIN_REG_LEN) {
		bytes = MIN_REG_LEN;
	}

	bytes_as_sword = (uns_sword)bytes + HDR_LEN;

	/* Make sure to check for header space when allocating */
	if (!enough_space(ctx, bytes_as_sword)) {
		uns_cheney_c89_collect(gc);
		if (!enough_space(ctx, bytes_as_sword)) {
			uns_on_oom(gc);
			return NULL;
		}
	}

	return raw_alloc(ctx, HDR_LEN + bytes, is_record);
}

void *uns_cheney_c89_alloc(Uns_GC gc, size_t bytes, enum uns_bytes_type typ)
{
	assert(typ == UNS_NOEXEC);
	return alloc(gc, bytes, 0);
}

void *uns_cheney_c89_alloc_rec(Uns_GC gc, size_t len, enum uns_record_type typ)
{
	void *p;
	size_t i;
	assert(typ == UNS_POINTERS_ONLY);

	if (len >= SIZE_MAX/sizeof(void*)) {
	    	uns_on_oom(gc);
		return NULL;
	} else if (len == 0) {
		len = 1;
	}

	p = alloc(gc, len*sizeof(void*), 1);
	if (p == NULL)
		return NULL;

	for (i = 0; i < len; i++)
		uns_cheney_c89_set(p, i, UNS_POINTER, NULL);
	return p;
}

int uns_cheney_c89_init(Uns_GC gc, size_t heap_size)
{
	struct ctx *ctx = malloc(sizeof(struct ctx));
	if (!ctx)
		return 0;

	uns_deinit(gc);
	ctx->tospace_alloc = ctx->tospace = malloc(heap_size);
	VALGRIND_CREATE_MEMPOOL(ctx->tospace, REDZONE, 0);
	if (!ctx->tospace) {
		free(ctx);
		return 0;
	}

	ctx->stats.usage_before = ctx->stats.usage_after
		= ctx->stats.collection_number = 0;

	ctx->tospace_end = ctx->tospace + heap_size;
	ctx->new_heap_size = heap_size;
	ctx->cb = NULL;

	uns_set_ctx(gc, ctx);
	uns_set_deinit(gc, uns_cheney_c89_deinit);
	uns_set_collect(gc, uns_cheney_c89_collect);
	uns_set_alloc(gc, uns_cheney_c89_alloc);
	uns_set_alloc_rec(gc, uns_cheney_c89_alloc_rec);
	uns_set_set(gc, uns_cheney_c89_set);
	uns_set_get(gc, uns_cheney_c89_get);

	return 1;
}