allow direct calls between compiled functions

3 files changed, 78 insertions, 27 deletions

diff --git a/src/common.h b/src/common.h
index 64ca250..5cc0c7f 100644
--- a/src/common.h
+++ b/src/common.h
@@ -257,6 +257,7 @@ struct ejit_func {
 	size_t fpr;
 
 	void *arena;
+	void *direct_call;
 	size_t size;
 };
 
diff --git a/src/compile/compile.c b/src/compile/compile.c
index 2d1ec5f..0411a3e 100644
--- a/src/compile/compile.c
+++ b/src/compile/compile.c
@@ -106,6 +106,7 @@ static jit_gpr_t getgpr(struct ejit_func *f, size_t l, size_t i)
 
 static jit_fpr_t getfpr(struct ejit_func *f, size_t l, size_t i)
 {
+	(void)(f);
 	if (l < jit_vf_num())
 		return jit_vf(l);
 
@@ -1610,28 +1611,72 @@ static void compile_imm_call(jit_state_t *j, struct operands *src, struct operan
 	jit_calli(j, addr, argc, args);
 
 	jit_shrink_stack(j, fixup);
-	operands_reset(src);
-	operands_reset(dst);
+}
+
+/* adds a header that converts from our external interface (params on stack) to
+ * an internal format that's effectively just the underlying ABI, makes calls
+ * between functions we know are compiled a bit faster. */
+static void compile_trampoline(struct ejit_func *f, jit_state_t *j)
+{
+	/* very important, argc we don't really do anything with but JIT_R1
+	 * contains the argument stack! */
+
+	size_t frame = jit_enter_jit_abi(j, 0, 0, 0);
+	jit_load_args_2(j,
+	                jit_operand_gpr(JIT_OPERAND_ABI_WORD, JIT_R0),
+	                jit_operand_gpr(JIT_OPERAND_ABI_POINTER, JIT_R1));
+
+	struct operands args = operands_create(0);
+
+	foreach_vec(ii, f->insns) {
+		struct ejit_insn i = *insns_at(&f->insns, ii);
+		switch (i.op) {
+		case PARAM: {
+			jit_operand_t p = jit_operand_mem(jit_abi_from(i.r1), JIT_R1, arg_offset(i));
+			operands_append(&args, p);
+			break;
+		}
+
+		case PARAM_F: {
+			jit_operand_t p = jit_operand_mem(jit_abi_from(i.r1), JIT_R1, arg_offset(i));
+			operands_append(&args, p);
+			break;
+		}
+
+		case START: {
+			/* callr should be smart enough to avoid JIT_R0 if some
+			 * other register wants to write to it */
+			jit_reloc_t r = jit_mov_addr(j, JIT_R0);
+			jit_callr(j, JIT_R0, operands_len(&args), args.buf);
+			jit_leave_jit_abi(j, 0, 0, frame);
+			jit_ret(j); /* should just forward the return value */
+
+			f->direct_call = jit_address(j);
+			jit_patch_here(j, r);
+
+			operands_destroy(&args);
+			return;
+		}
+
+		default: abort();
+		}
+	}
 }
 
 static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
                               size_t size)
 {
 	jit_begin(j, arena, size);
+	compile_trampoline(f, j);
+
 	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);
-
-	/* very important, argc we don't really do anything with but JIR_R1
-	 * contains the argument stack! */
-	jit_load_args_2(j,
-	                jit_operand_gpr(JIT_OPERAND_ABI_WORD, JIT_R0),
-	                jit_operand_gpr(JIT_OPERAND_ABI_POINTER, JIT_R1));
-
 	size_t stack = jit_align_stack(j, stack_size(f));
 
 	struct operands src = operands_create();
 	struct operands dst = operands_create();
+	struct operands direct = operands_create();
 
 	struct relocs relocs = relocs_create(sizeof(struct reloc_helper));
 	struct addrs addrs = addrs_create();
@@ -1846,6 +1891,7 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 
 			operands_append(&src, type);
 			operands_append(&src, arg);
+			operands_append(&direct, arg);
 
 			jit_operand_t to[2] = {
 				jit_operand_mem(JIT_OPERAND_ABI_WORD, JIT_SP,
@@ -1878,6 +1924,7 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 
 			operands_append(&src, type);
 			operands_append(&src, arg);
+			operands_append(&direct, arg);
 
 			jit_operand_t to[2] = {
 				jit_operand_mem(JIT_OPERAND_ABI_WORD, JIT_SP,
@@ -1898,6 +1945,9 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 				jit_operand_gpr(JIT_OPERAND_ABI_POINTER, JIT_SP)
 			};
 			compile_imm_call(j, &src, &dst, (void *)i.o, 2, args);
+			operands_reset(&src);
+			operands_reset(&dst);
+			operands_reset(&direct);
 			break;
 		}
 
@@ -1908,10 +1958,22 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 				jit_operand_gpr(JIT_OPERAND_ABI_POINTER, JIT_SP)
 			};
 			compile_imm_call(j, &src, &dst, (void *)i.o, 2, args);
+			operands_reset(&src);
+			operands_reset(&dst);
+			operands_reset(&direct);
 			break;
 		}
 
 		case CALLI: {
+			struct ejit_func *f = (struct ejit_func *)i.o;
+			if (f && f->direct_call) {
+				jit_calli(j, f->direct_call, operands_len(&direct), direct.buf);
+				operands_reset(&src);
+				operands_reset(&dst);
+				operands_reset(&direct);
+				break;
+			}
+
 			jit_operand_t args[3] = {
 				jit_operand_imm(JIT_OPERAND_ABI_POINTER, i.o),
 				jit_operand_imm(JIT_OPERAND_ABI_WORD,
@@ -1919,6 +1981,9 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 				jit_operand_gpr(JIT_OPERAND_ABI_POINTER, JIT_SP)
 			};
 			compile_imm_call(j, &src, &dst, ejit_run_func, 3, args);
+			operands_reset(&src);
+			operands_reset(&dst);
+			operands_reset(&direct);
 			break;
 		}
 
@@ -1971,13 +2036,6 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 		}
 
 		case PARAM_F: {
-			/* move from argument stack to location */
-			jit_operand_t from = jit_operand_mem(
-				jit_abi_from(i.r1),
-				JIT_R1,
-				arg_offset(i)
-				);
-
 			jit_operand_t to;
 			if (i.r0 < jit_vf_num()) {
 				/* regular register */
@@ -1990,19 +2048,11 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 				                     stack_loc_f(f, i.r2));
 			}
 
-			operands_append(&src, from);
 			operands_append(&dst, to);
 			break;
 		}
 
 		case PARAM: {
-			/* move from argument stack to location */
-			jit_operand_t from = jit_operand_mem(
-				jit_abi_from(i.r1),
-				JIT_R1,
-				arg_offset(i)
-				);
-
 			jit_operand_t to;
 			if (i.r0 < jit_v_num()) {
 				/* regular register */
@@ -2015,17 +2065,15 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 				                     stack_loc(i.r2));
 			}
 
-			operands_append(&src, from);
 			operands_append(&dst, to);
 			break;
 		}
 
 		case START: {
 			/* parameters should be done by now */
-			jit_move_operands(j, dst.buf, src.buf, operands_len(&src));
+			jit_load_args(j, operands_len(&dst), dst.buf);
 			/* reuse for arguments */
 			operands_reset(&dst);
-			operands_reset(&src);
 			break;
 		}
 
@@ -2044,6 +2092,7 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 
 	operands_destroy(&src);
 	operands_destroy(&dst);
+	operands_destroy(&direct);
 	relocs_destroy(&relocs);
 	addrs_destroy(&addrs);
 
diff --git a/src/ejit.c b/src/ejit.c
index cfbbbc1..7a31c5a 100644
--- a/src/ejit.c
+++ b/src/ejit.c
@@ -52,6 +52,7 @@ struct ejit_func *ejit_create_func(enum ejit_type rtype, size_t argc,
 	f->insns = insns_create();
 	f->labels = labels_create();
 	f->arena = NULL;
+	f->direct_call = NULL;
 	f->size = 0;
 
 	for (size_t i = 0; i < argc; ++i) {