1 files changed, 39 insertions, 14 deletions

diff --git a/src/compile/compile.c b/src/compile/compile.c
index bfcb12d..5432bc1 100644
--- a/src/compile/compile.c
+++ b/src/compile/compile.c
@@ -52,6 +52,19 @@ static void assert_helper(const char *msg)
 	assert(false && msg);
 }
 
+static bool gpr_free(size_t argc, jit_operand_t args[argc], jit_gpr_t r)
+{
+	for (size_t i = 0; i < argc; ++i) {
+		if (args[i].kind != JIT_OPERAND_KIND_GPR)
+			continue;
+
+		if (jit_gpr_regno(args[i].loc.gpr.gpr) == jit_gpr_regno(r))
+			return false;
+	}
+
+	return true;
+}
+
 static void free_arena(void *arena, size_t size)
 {
 	munmap(arena, size);
@@ -2489,7 +2502,7 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 				 * argument stack address */
 				jit_operand_gpr(JIT_OPERAND_ABI_POINTER, JIT_R0)
 			};
-			compile_imm_call(j, &src, &dst, (void *)(uintptr_t)i.o, 2, args);
+			compile_imm_call(j, &src, &dst, (void *)i.p, 2, args);
 			restore_caller_save_regs(f, j);
 
 			operands_reset(&src);
@@ -2502,7 +2515,7 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 			/* a bit of copy-paste between this and the next func,
 			 * hmm */
 			assert(operands_len(&direct) <= 2);
-			struct ejit_func *f = (struct ejit_func *)(uintptr_t)i.o;
+			struct ejit_func *f = (struct ejit_func *)i.p;
 			assert(f->direct_call);
 
 			jit_operand_t regs[2] = {
@@ -2560,32 +2573,44 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 						(jit_imm_t)"trying to tail call interpreted function"));
 			jit_patch_here(j, direct_reloc);
 #endif
+			size_t argc = operands_len(&direct);
 
+			/* r0 = target, r1 = arg1, r2 = arg2 */
 			jit_ldxi(j, JIT_R0, r, offsetof(struct ejit_func, direct_call));
-			jit_operand_t regs[2] = {
+			jit_operand_t regs[3] = {
 				jit_operand_gpr(JIT_OPERAND_ABI_WORD, JIT_R1),
 				jit_operand_gpr(JIT_OPERAND_ABI_WORD, JIT_R2)
 			};
-			jit_move_operands(j, regs, direct.buf, operands_len(&direct));
+			jit_move_operands(j, regs, direct.buf, argc);
 
 			/* with args safely in registers, reset stack/state
 			 * while avoiding overwriting the call target */
-			jit_gpr_t tmp = get_callr_temp(j);
-			jit_movr(j, tmp, JIT_R0);
-
 			int frame_size = j->frame_size;
 			jit_shrink_stack(j, stack);
 			jit_leave_jit_abi(j, gprs, fprs, frame);
 
-			/* now move args into place */
-			jit_operand_t args[2] = {};
-			foreach_vec(oi, direct) {
+			/* now move args into place, making sure we avoid our
+			 * target register  */
+			jit_operand_t args[3] = {};
+			for (size_t oi = 0; oi < argc; ++oi) {
 				args[oi] = *operands_at(&direct, oi);
 			}
 
-			jit_locate_args(j, operands_len(&direct), args);
-			jit_move_operands(j, args, regs, operands_len(&direct));
-			jit_jmpr(j, tmp);
+			jit_locate_args(j, argc, args);
+
+			/* we know that at least one gpr must be free */
+			jit_gpr_t target = gpr_free(argc, args, JIT_R0) ? JIT_R0
+				         : gpr_free(argc, args, JIT_R1) ? JIT_R1
+					 : gpr_free(argc, args, JIT_R2) ? JIT_R2
+					 : (abort(), JIT_R0);
+
+			/* move our target in JIT_R0 to whatever the free
+			 * register is to avoid it being clobbered when we move
+			 * the actual arguments */
+			args[argc] = jit_operand_gpr(JIT_OPERAND_ABI_POINTER, target);
+			regs[argc] = jit_operand_gpr(JIT_OPERAND_ABI_POINTER, JIT_R0);
+			jit_move_operands(j, args, regs, argc + 1);
+			jit_jmpr(j, target);
 			j->frame_size = frame_size;
 
 			operands_reset(&src);
@@ -2641,7 +2666,7 @@ static size_t compile_fn_body(struct ejit_func *f, jit_state_t *j, void *arena,
 		case EJIT_OP_CALLI: {
 			save_caller_save_regs(f, j);
 
-			struct ejit_func *f = (struct ejit_func *)(uintptr_t)i.o;
+			struct ejit_func *f = (struct ejit_func *)i.p;
 #if __WORDSIZE != 64
 			assert(f->rtype != EJIT_INT64 && f->rtype != EJIT_UINT64);
 #endif