/*
 * top.c
 * @copyright (C) 2021 SML# Development Team.
 * @author UENO Katsuhiro
 */

#include "smlsharp.h"
#include <stdlib.h>
#include <string.h>

/*
 * A frame table generated by smlsharp_gc plugin is an array of
 * descriptors of the following form (see also smlsharp_gc.cpp):
 *
 *   struct __align__(intptr_t) {
 *     uint16_t NumSafePoints;
 *     uint16_t FrameSize;                      // in words
 *     uint16_t NumRoots;
 *     uint16_t RootOffsets[NumRoots];          // in words
 *     intptr_t SafePoints[NumSafePoints];      // in bytes
 *   } Descriptors[]; // terminated with NumSafePoints == 0
 *
 * There is no padding between adjacent descriptors.
 * The last descriptor has NumSafePoints = 0.
 * The size of each descriptor depends on FrameSize and NumRoots.
 *
 * Each descriptor describes the following properties of a function:
 * 1) the size of the stack frame of the function (FrameSize),
 * 2) the address where root pointers are stored (RootOffsets), and
 * 3) the code address where gc may occur (SafePoints).
 *
 * Each element of RootOffsets is an offset (in words) from the the end address
 * of the stack frame.
 * Each element of SafePoints is an offset (in bytes) from sml_main, which is
 * the second argument of the sml_gcroot function.
 * Both the RootOffsets and SafePoints are not sorted.
 */
/* struct sml_frame_layout is defined in smlsharp.h */
#define LAYOUT_SIZE(layout) \
	(sizeof(struct sml_frame_layout) \
	 + sizeof(uint16_t) * (layout)->num_roots)
#define SAFE_POINTS_OFFSET(layout) \
	CEILING(LAYOUT_SIZE(layout), sizeof(intptr_t))
#define SAFE_POINTS(layout) \
	((const intptr_t *)((char*)(layout) + SAFE_POINTS_OFFSET(layout)))
#define NEXT_LAYOUT(layout) \
	((const struct sml_frame_layout *) \
	 (SAFE_POINTS(layout) + (layout)->num_safe_points))

/*
 * A toplevel mutable object array generated by LLVMGen is of the form
 * (see also LLVMGen.sml):
 *
 *   struct {
 *     uintptr_t NumObjects;
 *     intptr_t ObjectOffsets[NumObjects];
 *   } TopObjects;
 *
 * ObjectOffsets are offsets (in bytes) from the address of TopObjects
 * to toplevel objects.
 *
 * TOpObjects is placed in the DATA section so that it can be computed at
 * compile time, whereas it is actuallly an immutable object.
 */
struct toplevel_objects {
	uintptr_t num_objects;
	intptr_t object_offsets[];
};

struct sml_gcroot {
	_Atomic(struct sml_gcroot *) next;
	size_t num_points;  /* the number of safe_points in item[] */
	size_t num_tops;    /* the number of toplevel_objects in item[] */
	size_t allocsize;
	struct safe_point {
		/* addr is either a struct toplevel_objects (if layout is null)
		 * or address of a call instruction (otherwise). */
		const void *addr;
		const struct sml_frame_layout *layout;
	} item[];
	/* item[0 ... num_points] is for safe_points sorted by addr
	 * item[num_layouts ... (num_layouts + num_tops)] is for
	 * toplevel_objects */
};
#define REST_SAFE_POINTS(gcroot) \
	(((gcroot)->allocsize - sizeof(struct sml_gcroot)) \
	 / sizeof(struct safe_point) \
	 - ((gcroot)->num_points + (gcroot)->num_tops))
#define TOPS(gcroot) (&(gcroot)->item[(gcroot)->num_points])
#define PAGE_SIZE CEILING(sizeof(struct sml_gcroot), 4096)

static struct sml_gcroot *
new_gcroot()
{
	struct sml_gcroot *gcroot;
	gcroot = sml_xmalloc(PAGE_SIZE);
	atomic_init(&gcroot->next, NULL);
	gcroot->num_points = 0;
	gcroot->num_tops = 0;
	gcroot->allocsize = PAGE_SIZE;
	return gcroot;
}

static struct safe_point *
alloc_safe_points(struct sml_gcroot **gcroot, size_t inc)
{
	struct sml_gcroot *p = *gcroot;
	struct safe_point *ret;
	size_t minsize, allocsize;

	if (REST_SAFE_POINTS(p) < inc) {
		minsize = p->num_points + p->num_tops + inc;
		minsize = minsize * sizeof(struct safe_point);
		minsize = sizeof(struct sml_gcroot) + minsize;
		allocsize = CEILING(minsize, PAGE_SIZE);
		p = sml_xrealloc(p, allocsize);
		p->allocsize = allocsize;
		*gcroot = p;
	}
	ret = &p->item[p->num_points + p->num_tops];
	p->num_points += inc;
	return ret;
}

void
sml_gcroot(void *gcroot_p, void (*smltabb)(void), void *smlftab, void *smlroot)
{
	struct sml_gcroot **gcroot = gcroot_p;
	const struct sml_frame_layout *ftab = smlftab, *t;
	const struct toplevel_objects *tops = smlroot;
	struct safe_point *dst;
	const intptr_t *points;
	size_t inc = 0, i;

	if (smlftab) {
		for (t = ftab; t->num_safe_points > 0; t = NEXT_LAYOUT(t))
			inc += t->num_safe_points;
	}
	if (smlroot)
		inc += 1;
	dst = alloc_safe_points(gcroot, inc);

	if (smlftab) {
		for (t = ftab; t->num_safe_points > 0; t = NEXT_LAYOUT(t)) {
			points = SAFE_POINTS(t);
			for (i = 0; i < t->num_safe_points; i++) {
				dst->addr = (char*)smltabb + points[i];
				dst->layout = t;
				dst++;
			}
		}
	}
	if (smlroot) {
		dst->addr = tops;
		dst->layout = NULL;
		dst++;
		(*gcroot)->num_tops++;
		(*gcroot)->num_points--;
	}
}

static int
cmp_safe_point(const void *p1, const void *p2)
{
	const struct safe_point *s1 = p1, *s2 = p2;
	uintptr_t n1, n2;
	if (s1->layout && !s2->layout)
		return -1;
	if (!s1->layout && s2->layout)
		return 1;
	n1 = (uintptr_t)s1->addr;
	n2 = (uintptr_t)s2->addr;
	if (n1 > n2)
		return 1;
	if (n1 < n2)
		return -1;
	return 0;
}

static void
sort_safe_points(struct sml_gcroot *gcroot)
{
	qsort(gcroot->item,
	      gcroot->num_points + gcroot->num_tops,
	      sizeof(struct safe_point),
	      cmp_safe_point);
}

/* gcroot_list is maintained as follows:
 *
 * [first program] -> [last program] -> ... -> [second program]
 *
 * based on the observation that the first (main) and last programs are
 * frequently executed and therefore their safe points are frequently
 * searched for.
 */
static _Atomic(struct sml_gcroot *) gcroot_list = NULL;

static void
register_gcroot(struct sml_gcroot *gcroot)
{
	struct sml_gcroot *old, *first;

	first = load_relaxed(&gcroot_list);
	if (!first && cmpswap_release(&gcroot_list, &first, gcroot))
		return;
	do {
		old = load_relaxed(&first->next);
		store_relaxed(&gcroot->next, old);
	} while (!cmpswap_weak_release(&first->next, &old, gcroot));
}

struct sml_gcroot *
sml_gcroot_load(void (* const *sml_loads)(void *), unsigned int count)
{
	unsigned int i;

	struct sml_gcroot *gcroot;
	gcroot = new_gcroot();
	for (i = 0; i < count; i++)
		sml_loads[i](&gcroot);
	sort_safe_points(gcroot);
	register_gcroot(gcroot);
	return gcroot;
}

void
sml_gcroot_unload(struct sml_gcroot *gcroot)
{
	/* ToDo: free gcroot memory */
	gcroot->num_points = 0;
	gcroot->num_tops = 0;
}

static const struct safe_point *
binary_search(void *keyaddr, const struct safe_point *begin, size_t len)
{
	uintptr_t addr, key = (uintptr_t)keyaddr;
	size_t l = 0, r = len, c;

	while (l < r) {
		c = (l + r) / 2;
		addr = (uintptr_t)begin[c].addr;
		if (addr == key)
			return &begin[c];
		else if (addr < key)
			l = c + 1;
		else
			r = c;
	}
	return NULL;
}

const struct sml_frame_layout *
sml_lookup_frametable(void *retaddr)
{
	struct sml_gcroot *p;
	const struct safe_point *s;

	for (p = load_acquire(&gcroot_list); p; p = load_acquire(&p->next)) {
		s = binary_search(retaddr, p->item, p->num_points);
		if (s)
			return s->layout;
	}
	sml_fatal(0, "frametable lookup failed %p", retaddr);
}

void
sml_global_enum_ptr(void (*trace)(void **, void *), void *data)
{
	struct sml_gcroot *p;
	const struct safe_point *items;
	const struct toplevel_objects *tops;
	void *obj;
	size_t i, j;

	for (p = load_acquire(&gcroot_list); p; p = load_acquire(&p->next)) {
		items = TOPS(p);
		for (i = 0; i < p->num_tops; i++) {
			assert(items[i].layout == NULL);
			tops = items[i].addr;
			for (j = 0; j < tops->num_objects; j++) {
				obj = (char*)tops + tops->object_offsets[j];
				sml_obj_enum_ptr(obj, trace, data);
			}
		}
	}
}
