x86-64 o3: direct o3-to-o3 calls (rcall_x64 / rfinish_x64)

Port arm64's rcall/rfinish to the x86-64 region jit: call the callee's entry
stub through fn_nat with a built frame, return via linkstk, and fall back to
o3_invoke when the stub is missing or the frame/depth check fails. Clear
fn_nat on rforget and cover the function prologue so icon.replace_function
drops the region (o3_caller_replace). vm_until saves jit_depth around nested
bytecode, matching arm64.

1 files changed, 182 insertions, 14 deletions

diff --git a/src/core/core.c b/src/core/core.c
index e7b1ef6..748c508 100644
--- a/src/core/core.c
+++ b/src/core/core.c
@@ -3309,6 +3309,13 @@ static Rg *rat(int p)
 #ifdef ZYNATIVE
 static VM_TLS unsigned char *cov;	/* how many live translations cover the word (255 sticks) */
 static VM_TLS void **tab;		/* vm address -> translation */
+#if defined(ZYJIT) || defined(ZYJIT_X64)
+typedef struct { void *stub; int fslots; int argc; } FnNat;
+static VM_TLS FnNat *fn_nat;
+#define LINKN 1024
+static VM_TLS void *linkstk[LINKN];
+static VM_TLS Cell jit_depth;
+#endif
 static void be_drop(int j) { tab[ix_blk[j].a] = 0; cdec(cov, ix_blk[j].a, ix_blk[j].b); if (trace) fprintf(stderr, "jit: dropped %d..%d\n", ix_blk[j].a, ix_blk[j].b - 1); }
 #ifdef ZYJIT
 typedef unsigned JW;		/* code is a string of 32-bit instruction words */
@@ -3630,7 +3637,12 @@ static int hook(Cell a, Cell v)
 	m[a] = v;
 	hit = ix_drop_at((int)a);
 	for (j = 0; j < nrg; j++)
-		if (rgs[j]->lifted && rgs[j]->lo <= a && a < rgs[j]->hi) {
+		if (rgs[j]->lifted && ((rgs[j]->lo <= a && a < rgs[j]->hi) ||
+		    ((unsigned)rgs[j]->fn < (unsigned)code_words && a >= rgs[j]->fn && a < rgs[j]->fn + 2))) {
+			/* a store into the function prologue replaces the function: the region is done.
+			   a store into the body only (lo..hi) clears it so it can lift again. */
+			if ((unsigned)rgs[j]->fn < (unsigned)code_words && a >= rgs[j]->fn && a < rgs[j]->fn + 2)
+				rgs[j]->dead = 1;
 			if (trace) fprintf(stderr, "o3: dropped %d..%d (a store to %d)\n", rgs[j]->lo, rgs[j]->hi - 1, (int)a);
 			rclear(rgs[j]);
 			hit = 1;
@@ -3981,11 +3993,6 @@ static VM_TLS unsigned char rk[MAXUOP * 4 + 16], rhome[MAXUOP * 4 + 16], rfp[MAX
 static VM_TLS Cell rkv[MAXUOP * 4 + 16]; static VM_TLS int rlast[MAXUOP * 4 + 16];
 static VM_TLS unsigned **rfa; static VM_TLS int *rfb, rnf;	/* branches to blocks, fixed up at the end */
 
-#define LINKN 1024
-static VM_TLS void *linkstk[LINKN];
-static VM_TLS Cell jit_depth;
-typedef struct { void *stub; int fslots; int argc; } FnNat;
-static VM_TLS FnNat *fn_nat;
 static void rforget(Rg *R)
 {
 	int e;
@@ -5323,9 +5330,20 @@ static void rmark(Rg *R)
 {
 	int e;
 	R->marked = 1; cinc(rcov, R->lo, R->hi); cinc(cov, R->lo, R->hi);
+	/* cover the function prologue (R->fn): icon.replace_function patches it with
+	   lit+jmp, and a direct o3 caller must not keep the old stub when that happens. */
+	if ((unsigned)R->fn < (unsigned)code_words) {
+		cinc(rcov, R->fn, R->fn + 2); cinc(cov, R->fn, R->fn + 2);
+		tab[R->fn] = 0;
+		if ((unsigned)(R->fn + 1) < (unsigned)code_words) tab[R->fn + 1] = 0;
+	}
 	for (e = 0; e < R->ne; e++) if ((unsigned)R->ent[e] < (unsigned)code_words) tab[R->ent[e]] = 0;
 }
-static void runmark(Rg *R) { cdec(rcov, R->lo, R->hi); cdec(cov, R->lo, R->hi); }
+static void runmark(Rg *R)
+{
+	cdec(rcov, R->lo, R->hi); cdec(cov, R->lo, R->hi);
+	if ((unsigned)R->fn < (unsigned)code_words) { cdec(rcov, R->fn, R->fn + 2); cdec(cov, R->fn, R->fn + 2); }
+}
 #else
 static void rforget(Rg *R) { (void)R; }
 static void rmark(Rg *R) { R->marked = 1; cinc(rcov, R->lo, R->hi); }
@@ -6298,6 +6316,123 @@ static void rvset(Rg *R, int j, int r)
 }
 static VM_TLS JW **rfa; static VM_TLS int *rfb, nrf;
 static void rjmp(JW *at, int b) { rfa[nrf] = at; rfb[nrf++] = b; }
+/* leave a region to a native caller: restore the frame, drop depth, ret. depth 0
+   leaves through x_exit (bytecode entered the region). */
+static void rfinish_x64(Rg *R)
+{
+	JW *z, *ne;
+	mov_ri(r_cx, (Cell)(uintptr_t)&jit_depth);
+	mov_rm(r_ax, r_cx, 0);
+	test_rr(r_ax, r_ax);
+	z = j32(c_e);
+	mov_ri(r_cx, R->hi);
+	alu_rr(a_cmp, r_dx, r_cx);
+	ne = j32(c_ne);
+	alu_ri(a_sub, rSP, 8);
+	mov_mi(rSP, 0, 0);
+	fix32(ne, cp);
+	if (R->fslots) {
+		mov_rm(r_ax, rMB, 8 * ZFP);
+		alu_ri(a_add, r_ax, R->fslots);
+		mov_mr(rMB, 8 * ZFP, r_ax);
+	}
+	mov_ri(r_cx, (Cell)(uintptr_t)&jit_depth);
+	mov_rm(r_ax, r_cx, 0);
+	alu_ri(a_sub, r_ax, 1);
+	mov_mr(r_cx, 0, r_ax);
+	mov_ri(r_cx, (Cell)(uintptr_t)linkstk);
+	mov_rx(r_ax, r_cx, r_ax, 8, 0);
+	jmp_r(r_ax);
+	fix32(z, cp);
+	mov_rr(r_ax, r_dx);
+	jmp_to(x_exit);
+}
+/* o3-to-o3: call the callee stub with a built frame. else a vm call via x_disp.
+   generation guard: fn_nat stub is cleared in rforget when the callee drops. */
+static void rcall_x64(Rg *R, RB *B)
+{
+	JW *to_slow, *is_null, *too_deep, *too_small, *argc0, *loop;
+	int e, hopped;
+	rspill(R);
+	/* rdx = callee target */
+	mov_ri(r_cx, (Cell)(uintptr_t)&code_words);
+	alu_rm(a_cmp, r_dx, r_cx, 0);
+	to_slow = j32(c_ae);
+	mov_ri(r_ax, (Cell)(uintptr_t)&fn_nat);
+	mov_rm(r_ax, r_ax, 0);
+	mov_rr(r_cx, r_dx);
+	shift_ri(4, r_cx, 4);
+	alu_rr(a_add, r_ax, r_cx);
+	mov_rm(r_9, r_ax, 0);
+	test_rr(r_9, r_9);
+	is_null = j32(c_e);
+	mov_ri(r_cx, (Cell)(uintptr_t)&jit_depth);
+	mov_rm(r_8, r_cx, 0);
+	alu_ri(a_cmp, r_8, LINKN);
+	too_deep = j32(c_ae);
+	mov_r32m(r_si, r_ax, 8);
+	mov_r32m(r_di, r_ax, 12);
+	mov_rm(r_ax, rMB, 8 * ZFP);
+	alu_rr(a_sub, r_ax, r_si);
+	mov_ri(r_cx, (Cell)(uintptr_t)&frame_floor_word);
+	mov_rm(r_cx, r_cx, 0);
+	mov_rx(r_8, rMB, r_cx, 8, 0);
+	mov_ri(r_cx, (Cell)(uintptr_t)&frame_guard_word);
+	mov_rm(r_cx, r_cx, 0);
+	mov_rx(r_cx, rMB, r_cx, 8, 0);
+	alu_rr(a_add, r_8, r_cx);
+	alu_rr(a_cmp, r_ax, r_8);
+	too_small = j32(c_b);
+	push_r(r_9);
+	mov_mr(rMB, 8 * ZFP, r_ax);
+	mov_rr(r_8, r_ax);
+	test32(r_di, r_di);
+	argc0 = j32(c_e);
+	loop = cp;
+	alu_ri(a_sub, r_di, 1);
+	mov_rm(r_cx, rSP, 0);
+	alu_ri(a_add, rSP, 8);
+	lea(r_ax, r_8, r_di, 1, 0);
+	mov_xr(rMB, r_ax, 8, 0, r_cx);
+	test32(r_di, r_di);
+	jcc_to(c_ne, loop);
+	fix32(argc0, cp);
+	pop_r(r_9);
+	mov_ri(r_cx, (Cell)(uintptr_t)&jit_depth);
+	mov_rm(r_ax, r_cx, 0);
+	alu_ri(a_add, r_ax, 1);
+	mov_mr(r_cx, 0, r_ax);
+	alu_ri(a_sub, r_sp, 8);
+	call_r(r_9);
+	alu_ri(a_add, r_sp, 8);
+	hopped = 0;
+	for (e = 0; e < R->ne; e++) if (R->ent[e] == B->x2) { jmp_to((JW *)R->stub[e]); hopped = 1; break; }
+	if (!hopped) { mov_r32i(r_ax, (unsigned)B->x2); jmp_to(x_disp); }
+	fix32(to_slow, cp);
+	fix32(is_null, cp);
+	fix32(too_deep, cp);
+	fix32(too_small, cp);
+	/* slow: publish stacks (region code: no vd), o3_invoke, resume */
+	mov_rr(r_cx, rSP); alu_rr(a_sub, r_cx, rMB); shift_ri(5, r_cx, 3); mov_mr(rST, 8, r_cx);
+	mov_rr(r_cx, rRP); alu_rr(a_sub, r_cx, rMB); shift_ri(5, r_cx, 3); mov_mr(rST, 16, r_cx);
+	lea(r_di, rST, -1, 1, 8);
+	lea(r_si, rST, -1, 1, 16);
+	/* rdx still holds target from the early path */
+	mov_ri(r_cx, B->x2);
+	call_abs((void *)o3_invoke);
+	reload_stacks();
+	mov_ri(r_cx, (Cell)(uintptr_t)&halt_at);
+	mov_rm(r_ax, r_cx, 0);
+	alu_ri(a_cmp, r_ax, 0);
+	{
+		JW *ok = j32(c_l);
+		jmp_to(x_exit);
+		fix32(ok, cp);
+	}
+	hopped = 0;
+	for (e = 0; e < R->ne; e++) if (R->ent[e] == B->x2) { jmp_to((JW *)R->stub[e]); hopped = 1; break; }
+	if (!hopped) { mov_r32i(r_ax, (unsigned)B->x2); jmp_to(x_disp); }
+}
 static void remit_x64(Rg *R, RB *B)
 {
 	int i, j, d;
@@ -6344,16 +6479,22 @@ static void remit_x64(Rg *R, RB *B)
 		if (B->s2 >= 0) rjmp(jm32(), B->s2); else rexit(R, B->x2);
 		break;
 	case T_CALL:
-		rspill(R);
-		mov_mi(rRP, 0, B->x2); alu_ri(a_add, rRP, 8);
-		mov_rr(r_ax, r_dx); jmp_to(x_disp);
+		rcall_x64(R, B);
 		break;
 	case T_RET:
 		rspill(R);
 		alu_ri(a_sub, rRP, 8); mov_rm(r_ax, rRP, 0); jmp_to(x_disp);
 		break;
-	default:
-		rspill(R); mov_rr(r_ax, r_dx); jmp_to(x_disp);
+	default: {
+		JW *hs;
+		rspill(R);
+		mov_ri(r_cx, R->hi);
+		alu_rr(a_cmp, r_dx, r_cx);
+		hs = j32(c_ae);
+		mov_rr(r_ax, r_dx); jmp_to(x_disp);
+		fix32(hs, cp);
+		rfinish_x64(R);
+		break; }
 	}
 }
 /* 0: compiled; -1: can't be; -2: no room now */
@@ -6379,7 +6520,19 @@ static int rjit_x64(Rg *R)
 	for (b = 0; b < R->np; b++) fprintf(stderr, "ps %d -> var %d\n", R->ps[b], R->pv[b]);
 #endif
 	for (e = 0; e < R->ne && !rc; e++) {
+		JW *skip;
 		R->stub[e] = cp;
+		if (e == 0) {
+			mov_ri(r_cx, (Cell)(uintptr_t)&jit_depth);
+			mov_rm(r_ax, r_cx, 0);
+			test_rr(r_ax, r_ax);
+			skip = j32(c_e);
+			alu_ri(a_sub, r_ax, 1);
+			pop_r(r_cx);
+			mov_ri(r_8, (Cell)(uintptr_t)linkstk);
+			mov_xr(r_8, r_ax, 8, 0, r_cx);
+			fix32(skip, cp);
+		}
 		mov_rm(r_ax, rMB, 8 * ZFP); lea(r_bp, rMB, r_ax, 8, 0);
 		rjmp(jm32(), R->entb[e]);
 	}
@@ -6394,6 +6547,11 @@ static int rjit_x64(Rg *R)
 		for (k = 0; k < nrf; k++) fix32(rfa[k], lab[rfb[k]]);
 		R->nat = start;
 		for (e = 0; e < R->ne; e++) if ((unsigned)R->ent[e] < (unsigned)code_words) tab[R->ent[e]] = R->stub[e];
+		if ((unsigned)R->fn < (unsigned)code_words) {
+			fn_nat[R->fn].stub = R->stub[0];
+			fn_nat[R->fn].fslots = R->fslots;
+			fn_nat[R->fn].argc = R->argc;
+		}
 		if (trace) fprintf(stderr, "o3: %d..%d -> %d bytes of x86-64\n", R->lo, R->hi - 1, (int)(cp - start));
 	} else cp = start;
 	free(lab); free(rfa); free(rfb);
@@ -6403,6 +6561,7 @@ static void rforget(Rg *R)
 {
 	int e;
 	if (R->nat) for (e = 0; e < R->ne; e++) if ((unsigned)R->ent[e] < (unsigned)code_words && tab[R->ent[e]] == R->stub[e]) tab[R->ent[e]] = 0;
+	if ((unsigned)R->fn < (unsigned)code_words) fn_nat[R->fn].stub = 0;
 	R->nat = 0;
 }
 static void rcompile(Rg *R, int eager)
@@ -6603,6 +6762,13 @@ static int vm_until(int pc, int stop_eq, int stop_ge, Rg *re, Cell *ps, Cell *pr
 	void *saved_eq = 0, *saved_ent[RMAXE];
 	int i, nent = 0, flushed = 0, ep0 = drop_epoch, term = 0;
 	Cell return_depth = *pr - (stack_control ? stack_get(SC_RETURN_LOW) : RS0) - 1;
+	Cell saved_depth = jit_depth, *saved_links = 0;
+	if (saved_depth) {
+		saved_links = malloc((size_t)saved_depth * sizeof *saved_links);
+		if (!saved_links) return PC_FRAME_OVERFLOW;
+		memcpy(saved_links, linkstk, (size_t)saved_depth * sizeof *saved_links);
+	}
+	jit_depth = 0;
 	enter.u = x_enter;
 	if (stop_eq >= 0 && stop_eq < code_words) { saved_eq = tab[stop_eq]; tab[stop_eq] = 0; }
 	if (re) {
@@ -6654,6 +6820,8 @@ static int vm_until(int pc, int stop_eq, int stop_ge, Rg *re, Cell *ps, Cell *pr
 		if (need_flush) { need_flush = 0; jitwrite(1); flush(); jitwrite(0); }
 	}
 	*ps = S[1]; *pr = S[2];
+	jit_depth = saved_depth;
+	if (saved_links) { memcpy(linkstk, saved_links, (size_t)saved_depth * sizeof *saved_links); free(saved_links); }
 	flushed = drop_epoch != ep0;
 	if (!flushed) {
 		if (stop_eq >= 0 && stop_eq < code_words) tab[stop_eq] = saved_eq;
@@ -6912,7 +7080,7 @@ static int code_reserve(Cell end)
 	CODE_COLUMN(o3ent, code_words, cap); CODE_COLUMN(o3ei, code_words, cap);
 #ifdef ZYNATIVE
 	CODE_COLUMN(tab, code_words, cap);
-#ifdef ZYJIT
+#if defined(ZYJIT) || defined(ZYJIT_X64)
 	CODE_COLUMN(fn_nat, code_words, cap);
 #endif
 #elif defined(ZYLLVM)
@@ -7004,7 +7172,7 @@ static void context_engine_free(void) {
  for (i = 0; i < NENH; i++) { free(enh[i].u); free(enh[i].reg); }
  free(pg_head); free(o3ent); free(o3ei); free(gn_); free(gh_);
 #ifdef ZYNATIVE
-#ifdef ZYJIT
+#if defined(ZYJIT) || defined(ZYJIT_X64)
  free(fn_nat);
 #endif
  free(tab); free(cov); if (code && code != MAP_FAILED && munmap(code, CODEW * 4) == 0) atomic_fetch_sub(&context_code_live, 1);