example compiles under jit

6 files changed, 207 insertions, 11 deletions

diff --git a/deps/lightening b/deps/lightening
-Subproject cd0824122864613ae58c07393de599ff9460c35
+Subproject 2d4e621bd03075b2db0b7fb8493a22f39d213c6
diff --git a/examples/main.c b/examples/main.c
index d33ed10..1677023 100644
--- a/examples/main.c
+++ b/examples/main.c
@@ -20,7 +20,7 @@ int main()
 
 	ejit_patch(f, r, l);
 
-	/* the highest location we used was 3, so we need to request 3 location
+	/* the highest location we used was 3, so we need to request 4 locations
 	 * for general purpose registers in total. No floating point registers,
 	 * so 0. */
 	ejit_compile_func(f, 4, 0);
diff --git a/src/common.h b/src/common.h
index e93e920..4f4fcd7 100644
--- a/src/common.h
+++ b/src/common.h
@@ -84,6 +84,7 @@ enum ejit_opcode {
 
 	START,
 	END,
+	LABEL,
 
 	OPCODE_COUNT,
 };
diff --git a/src/compile/compile.c b/src/compile/compile.c
index 513e0e1..e651199 100644
--- a/src/compile/compile.c
+++ b/src/compile/compile.c
@@ -2,7 +2,190 @@
 #include "../../deps/lightening/lightening/lightening.c"
 #include "../common.h"
 
+static void *alloc_arena(size_t size)
+{
+	return mmap(NULL, size,
+			PROT_EXEC | PROT_READ | PROT_WRITE,
+			MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+}
+
+static void free_arena(void *arena, size_t size)
+{
+	munmap(arena, size);
+}
+
+static size_t grploc_count(struct ejit_func *f)
+{
+	return f->gpr >= jit_v_num() ? 0 : f->gpr - jit_v_num();
+}
+
+static size_t frploc_count(struct ejit_func *f)
+{
+	return f->fpr >= jit_vf_num() ? 0 : f->fpr - jit_vf_num();
+}
+
+static size_t stack_size(struct ejit_func *f)
+{
+	return grploc_count(f) * sizeof(jit_uword_t)
+		+ frploc_count(f) * sizeof(jit_float64_t);
+}
+
+static jit_off_t stack_loc(struct ejit_func *f, size_t l)
+{
+	assert(l >= jit_v_num());
+	return (l - jit_v_num()) * sizeof(jit_uword_t);
+}
+
+static jit_off_t stack_loc_f(struct ejit_func *f, size_t l)
+{
+	assert(l >= jit_vf_num());
+	return grploc_count(f) * sizeof(jit_uword_t)
+		+ (l - jit_vf_num()) * sizeof(jit_float64_t);
+}
+
+struct reloc_helper {
+	jit_reloc_t r;
+	size_t to;
+};
+
+static jit_gpr_t getloc(struct ejit_func *f, jit_state_t *j, size_t l, size_t i)
+{
+	if (l < jit_v_num())
+		return jit_v(l);
+
+	jit_ldxi(j, jit_r(i), JIT_SP, stack_loc(f, l));
+	return jit_r(i);
+}
+
+static jit_gpr_t getreg(struct ejit_func *f, size_t l, size_t i)
+{
+	if (l < jit_v_num())
+		return jit_v(l);
+
+	return jit_r(i);
+}
+
+static void putloc(struct ejit_func *f, jit_state_t *j, size_t l, jit_gpr_t r)
+{
+	if (l < jit_v_num()) {
+		assert(jit_v(l).regno == r.regno);
+		return;
+	}
+
+	jit_stxi(j, stack_loc(f, l), JIT_SP, r);
+}
+
+static void compile_label(jit_state_t *j, size_t ii, struct vec *labels)
+{
+	vect_at(jit_addr_t, *labels, ii) = jit_address(j);
+}
+
+static void compile_movi(struct ejit_func *f, jit_state_t *j, struct ejit_insn i)
+{
+	jit_gpr_t r = getreg(f, i.r0, 0);
+	jit_movi(j, r, i.o);
+	putloc(f, j, i.r0, r);
+}
+
+static void compile_addr(struct ejit_func *f, jit_state_t *j, struct ejit_insn i)
+{
+	jit_gpr_t dst = getreg(f, i.r0, 0);
+	jit_gpr_t src0 = getloc(f, j, i.r1, 1);
+	jit_gpr_t src1 = getloc(f, j, i.r2, 2);
+	jit_addr(j, dst, src0, src1);
+	putloc(f, j, i.r0, dst);
+}
+
+static void compile_bltr(struct ejit_func *f, jit_state_t *j, struct ejit_insn i, struct vec *relocs)
+{
+	jit_gpr_t c0 = getloc(f, j, i.r0, 0);
+	jit_gpr_t c1 = getloc(f, j, i.r1, 1);
+	jit_reloc_t r = jit_bltr(j, c0, c1);
+	struct reloc_helper h = {.r = r, .to = i.o};
+	vect_append(struct reloc_helper, *relocs, &h);
+}
+
+static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena, size_t size)
+{
+	jit_begin(j, arena, size);
+	size_t gprs = f->gpr >= jit_v_num() ? jit_v_num() : f->gpr;
+	size_t fprs = f->fpr >= jit_vf_num() ? jit_vf_num() : f->fpr;
+	size_t frame = jit_enter_jit_abi(j, gprs, fprs, 0);
+
+	size_t stack = jit_align_stack(j, stack_size(f));
+
+	struct vec relocs = vec_create(sizeof(struct reloc_helper));
+	struct vec labels = vec_create(sizeof(jit_addr_t));
+	vec_reserve(&labels, vec_len(&f->insns));
+
+	foreach_vec(ii, f->insns) {
+		struct ejit_insn i = vect_at(struct ejit_insn, f->insns, ii);
+		switch (i.op) {
+		case MOVI: compile_movi(f, j, i); break;
+		case ADDR: compile_addr(f, j, i); break;
+
+		case BLTR: compile_bltr(f, j, i, &relocs); break;
+
+		case LABEL: compile_label(j, ii, &labels); break;
+		case RET: {
+			jit_gpr_t r = getloc(f, j, i.r0, 0);
+			/* R0 won't get overwritten by jit_leave_jit_abi */
+			jit_movr(j, JIT_R0, r);
+			jit_shrink_stack(j, stack);
+			jit_leave_jit_abi(j, gprs, fprs, frame);
+			jit_retr(j, JIT_R0);
+			break;
+		}
+
+		case START: continue;
+		case END: continue;
+		default: abort();
+		}
+	}
+
+	foreach_vec(ri, relocs) {
+		struct reloc_helper h = vect_at(struct reloc_helper, relocs, ri);
+		jit_addr_t a = vect_at(jit_addr_t, labels, h.to);
+		jit_reloc_t r = h.r;
+
+		assert(a);
+		jit_patch_there(j, r, a);
+	}
+
+	vec_destroy(&relocs);
+	vec_destroy(&labels);
+
+	if (jit_end(j, &size))
+		return 0;
+
+	return size;
+}
+
 bool ejit_compile(struct ejit_func *f)
 {
-	return false;
+	if (!init_jit())
+		return false;
+
+	jit_state_t *j = jit_new_state(NULL, NULL);
+	assert(j);
+
+	void *arena = NULL;
+	size_t size = 4096;
+
+	while (1) {
+		arena = alloc_arena(size);
+		assert(arena);
+
+		size_t required_size = compile_fn_body(f, j, arena, size);
+		if (required_size == 0)
+			break;
+
+		free_arena(arena, size);
+		size = required_size + 4096;
+	}
+
+	jit_destroy_state(j);
+	f->arena = arena;
+	f->size = size;
+	return true;
 }
diff --git a/src/ejit.c b/src/ejit.c
index 7fe2095..69ade9a 100644
--- a/src/ejit.c
+++ b/src/ejit.c
@@ -40,16 +40,15 @@ struct ejit_func *ejit_create_func(enum ejit_type rtype, size_t argc, const stru
 	f->arena = NULL;
 	f->size = 0;
 
-	emit_insn_i(f, START, 0, 0, 0);
-
 	for (size_t i = 0; i < argc; ++i) {
 		switch (args[i].kind) {
-		case EJIT_OPERAND_GPR: emit_insn_i(f, PARAM, args[i].r, 0, i); break;
-		case EJIT_OPERAND_FPR: emit_insn_i(f, PARAM_F, args[i].r, 0, i); break;
+		case EJIT_OPERAND_GPR: emit_insn_i(f, PARAM, args[i].r, 0, args[i].type); break;
+		case EJIT_OPERAND_FPR: emit_insn_i(f, PARAM_F, args[i].r, 0, args[i].type); break;
 		default: abort();
 		}
 	}
 
+	emit_insn_i(f, START, 0, 0, 0);
 	return f;
 }
 
@@ -58,6 +57,7 @@ void ejit_compile_func(struct ejit_func *f, size_t gpr, size_t fpr)
 	f->gpr = gpr;
 	f->fpr = fpr;
 
+	/* try to jit compile if possible */
 	if (ejit_compile(f))
 		return;
 
@@ -84,7 +84,9 @@ void ejit_destroy_func(struct ejit_func *f)
 
 struct ejit_label ejit_label(struct ejit_func *f)
 {
-	return (struct ejit_label){.addr = vec_len(&f->insns)};
+	size_t addr = vec_len(&f->insns);
+	emit_insn_r(f, LABEL, 0, 0, 0);
+	return (struct ejit_label){.addr = addr};
 }
 
 void ejit_patch(struct ejit_func *f, struct ejit_reloc r, struct ejit_label l)
@@ -187,9 +189,9 @@ void ejit_movi(struct ejit_func *s, struct ejit_gpr r0, long o)
 
 struct ejit_reloc ejit_bltr(struct ejit_func *s, struct ejit_gpr r0, struct ejit_gpr r1)
 {
-	struct ejit_label label = ejit_label(s);
+	size_t addr = vec_len(&s->insns);
 	emit_insn_i(s, BLTR, r0.r, r1.r, 0);
-	return (struct ejit_reloc){.insn = label.addr};
+	return (struct ejit_reloc){.insn = addr};
 }
 
 long ejit_run_func(struct ejit_func *f, size_t argc, struct ejit_arg args[argc])
@@ -197,7 +199,7 @@ long ejit_run_func(struct ejit_func *f, size_t argc, struct ejit_arg args[argc])
 	assert(f->gpr && "trying to run a function that hasn't been compiled");
 	assert(f->rtype == EJIT_VOID || ejit_int_type(f->rtype));
 	if (f->arena)
-		return ((ejit_escape_f_t)f)(argc, args);
+		return ((ejit_escape_t)f->arena)(argc, args);
 
 	return ejit_interp(f, argc, args, true, NULL).r;
 }
@@ -207,7 +209,7 @@ double ejit_run_func_f(struct ejit_func *f, size_t argc, struct ejit_arg args[ar
 	assert(f->fpr && "trying to run a function that hasn't been compiled");
 	assert(ejit_float_type(f->rtype));
 	if (f->arena)
-		return ((ejit_escape_f_t)f)(argc, args);
+		return ((ejit_escape_f_t)f->arena)(argc, args);
 
 	return ejit_interp(f, argc, args, true, NULL).d;
 }
diff --git a/src/vec.h b/src/vec.h
index 7a1d5cb..e998e27 100644
--- a/src/vec.h
+++ b/src/vec.h
@@ -92,4 +92,14 @@ static inline void vec_sort(struct vec *v, vec_comp_t comp)
 	qsort(v->buf, v->n, v->ns, (__compar_fn_t)comp);
 }
 
+static inline void vec_reserve(struct vec *v, size_t n)
+{
+	if (v->n >= n)
+		return;
+
+	v->n = n;
+	v->s = n;
+	v->buf = realloc(v->buf, v->s * v->ns);
+}
+
 #endif /* VEC_H */