/* zygote core: puts the genome (the zc compiler and its runtime, both vm code)
   into memory and runs it. on arm64, every block of vm code is translated to
   arm64 the first time it is reached, straight from live memory. stores that
   land on translated words throw those translations away, so code that
   rewrites itself, including the resident compiler, just gets retranslated
   on the next visit. on other hosts (or with -DZYINTERP) a plain interpreter
   with the same semantics runs instead. ZYTRACE=1 logs translations on stderr. */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#if defined(__aarch64__) && !defined(ZYINTERP)
#define ZYJIT 1
#include <sys/mman.h>
#ifdef __APPLE__
#include <pthread.h>
#include <libkern/OSCacheControl.h>
#endif
#endif
#include "genome.h"

#ifndef JTYPE
#define JTYPE char
#endif
char Q(void);
void G(char);
extern JTYPE J;

#define MW (1 << 20)	/* memory, in words */
#define DS0 (MW - 16)	/* data stack grows down from here */
#define RS0 0xF0000	/* return stack grows up from here */
#define CODEW (4 << 20)	/* 16 MB of code, in instruction words */
#define MAXB 65536
#define BLKW 1024	/* longest block, in vm words */

static int m[MW];
static int trace;
#ifdef ZYJIT
static unsigned char cov[MW];	/* 1 if a live translation covers the word */
static void *tab[MW];		/* vm address -> translation */
static struct { int a, b; } bl[MAXB];	/* live translations cover [a,b); a < 0 is dead */
static int nbl;
static unsigned *code, *cp, *x_enter, *x_exit, *x_disp, *x_first;
static long S[9];	/* m, s, r, tab, cov, Q, G, &J, store hook; read by the stubs */

enum { W0 = 0, W1 = 1, X9 = 9, X10, X11, M = 19, SP = 20, RP = 21, TAB = 22, ST = 23 };

static void E(unsigned x) { *cp++ = x; }
static void li(int d, int v)
{
	unsigned u = (unsigned)v;
	E(0x52800000 | (u & 0xffff) << 5 | d);
	if (u >> 16) E(0x72a00000 | (u >> 16) << 5 | d);
}
static void ldr_post(int t, int n, int i) { E(0xb8400400 | (i & 0x1ff) << 12 | n << 5 | t); }
static void ldr_pre(int t, int n, int i) { E(0xb8400c00 | (i & 0x1ff) << 12 | n << 5 | t); }
static void str_pre(int t, int n, int i) { E(0xb8000c00 | (i & 0x1ff) << 12 | n << 5 | t); }
static void str_post(int t, int n, int i) { E(0xb8000400 | (i & 0x1ff) << 12 | n << 5 | t); }
static void ldr_u(int t, int n, int o) { E(0xb9400000 | (o / 4) << 10 | n << 5 | t); }
static void str_u(int t, int n, int o) { E(0xb9000000 | (o / 4) << 10 | n << 5 | t); }
static void ldrx_u(int t, int n, int o) { E(0xf9400000 | (o / 8) << 10 | n << 5 | t); }
static void ldr_rr(int t, int n, int r, int ext) { E(0xb8600800 | r << 16 | ext << 13 | 1 << 12 | n << 5 | t); }
static void str_rr(int t, int n, int r, int ext) { E(0xb8200800 | r << 16 | ext << 13 | 1 << 12 | n << 5 | t); }
static void ldp_post(int t, int u, int n, int i) { E(0x28c00000 | ((i / 4) & 0x7f) << 15 | u << 10 | n << 5 | t); }
static void op3(unsigned op, int d, int n, int r) { E(op | r << 16 | n << 5 | d); }
static void addx_ext(int d, int n, int r, int ext) { E(0x8b200000 | r << 16 | ext << 13 | 2 << 10 | n << 5 | d); }
static void movw(int d, int r) { E(0x2a0003e0 | r << 16 | d); }
static void blr(int n) { E(0xd63f0000 | n << 5); }
static void b_to(unsigned *t) { long o = t - cp; E(0x14000000 | (o & 0x3ffffff)); }
static void cb_fix(unsigned *at, unsigned *to) { *at |= ((to - at) & 0x7ffff) << 5; }
#define SXTW 6
#define UXTW 2

static void build_stubs(void)
{
	cp = code;
	x_enter = cp;			/* int enter(long *S, void *code) */
	E(0xa9bc7bfd);			/* stp x29, x30, [sp, #-64]! */
	E(0x910003fd);			/* mov x29, sp */
	E(0xa90153f3);			/* stp x19, x20, [sp, #16] */
	E(0xa9025bf5);			/* stp x21, x22, [sp, #32] */
	E(0xf9001bf7);			/* str x23, [sp, #48] */
	E(0xaa0003f7);			/* mov x23, x0 */
	ldrx_u(M, ST, 0);
	ldr_u(X9, ST, 8); addx_ext(SP, M, X9, UXTW);
	ldr_u(X9, ST, 16); addx_ext(RP, M, X9, UXTW);
	ldrx_u(TAB, ST, 24);
	E(0xd61f0020);			/* br x1 */
	x_exit = cp;			/* w0 = next vm address */
	op3(0xcb000000, X9, SP, M); E(0xd342fd29); str_u(X9, ST, 8);	/* s = (x20 - x19) >> 2 */
	op3(0xcb000000, X9, RP, M); E(0xd342fd29); str_u(X9, ST, 16);
	E(0xa94153f3);			/* ldp x19, x20, [sp, #16] */
	E(0xa9425bf5);			/* ldp x21, x22, [sp, #32] */
	E(0xf9401bf7);			/* ldr x23, [sp, #48] */
	E(0xa8c47bfd);			/* ldp x29, x30, [sp], #64 */
	E(0xd65f03c0);			/* ret */
	x_disp = cp;			/* w0 = target: chain through tab[] or fall back to c */
	E(0x7144001f);			/* cmp w0, #MW */
	E(0x54000002 | ((x_exit - cp) & 0x7ffff) << 5);	/* b.hs exit */
	E(0xf8605ac9);			/* ldr x9, [x22, w0, uxtw #3] */
	E(0xb4000009 | ((x_exit - cp) & 0x7ffff) << 5);	/* cbz x9, exit */
	E(0xd61f0120);			/* br x9 */
	x_first = cp;
}

static void flush(void)
{
	memset(tab, 0, sizeof tab);
	memset(cov, 0, sizeof cov);
	nbl = 0;
	cp = x_first;
	if (trace) fprintf(stderr, "jit: flushed\n");
}

static void jitwrite(int on)
{
#ifdef __APPLE__
	pthread_jit_write_protect_np(!on);
#else
	(void)on;
#endif
}

static void syncicache(unsigned *a, unsigned *b)
{
#ifdef __APPLE__
	sys_icache_invalidate(a, (char *)b - (char *)a);
#else
	__builtin___clear_cache((char *)a, (char *)b);
#endif
}

/* store hook for words that are covered by a translation */
static int hook(int a, int v)
{
	int j, x, hit = 0;
	if (m[a] == v) return 0;
	m[a] = v;
	for (j = 0; j < nbl; j++)
		if (bl[j].a >= 0 && bl[j].a <= a && a < bl[j].b) {
			tab[bl[j].a] = 0;
			for (x = bl[j].a; x < bl[j].b; x++) cov[x] = 0;
			if (trace) fprintf(stderr, "jit: dropped %d..%d\n", bl[j].a, bl[j].b - 1);
			bl[j].a = -1;
			hit = 1;
		}
	for (j = 0; j < nbl; j++)
		if (bl[j].a >= 0) for (x = bl[j].a; x < bl[j].b; x++) cov[x] = 1;
	return hit;
}

static void op(int i, int k)	/* everything except control flow */
{
	unsigned *a, *b, *c, *z;
	switch (k) {
	case 1: ldr_u(X9, SP, 0); ldr_rr(X9, M, X9, SXTW); str_u(X9, SP, 0); return;
	case 2:	/* store: inline unless the target word is translated */
		ldp_post(X9, X10, SP, 8);
		E(0x7144013f);			/* cmp w9, #MW */
		a = cp; E(0x54000002);		/* b.hs done */
		ldrx_u(X11, ST, 32);
		E(0x3869496b);			/* ldrb w11, [x11, w9, uxtw] */
		b = cp; E(0x3500000b);		/* cbnz w11, slow */
		str_rr(X10, M, X9, UXTW);
		c = cp; E(0x14000000);		/* b done */
		cb_fix(b, cp);			/* slow: hook(a, v) */
		movw(W0, X9); movw(W1, X10); ldrx_u(X11, ST, 64); blr(X11);
		z = cp; E(0x34000000);		/* cbz w0, done */
		li(W0, i + 1); b_to(x_exit);	/* a translation died: resume at i+1 in c */
		cb_fix(a, cp); cb_fix(z, cp);
		*c |= (cp - c) & 0x3ffffff;
		return;
	case 3: ldr_u(X9, SP, 0); E(0x9100128a); ldr_rr(X9, X10, X9, SXTW); str_u(X9, SP, 0); return; /* add x10, x20, #4 */
	case 4: E(0x91001294); return;	/* add x20, x20, #4 */
	case 5: E(0x29402a89); E(0x2900268a); return;	/* ldp w9, w10, [x20]; stp w10, w9, [x20] */
	case 10: ldrx_u(X11, ST, 40); blr(X11); E(0x13001c09); str_pre(X9, SP, -4); return; /* sxtb w9, w0 */
	case 11:
		ldrx_u(X9, ST, 56);
		if (sizeof J == 1) E(0x39400129); else ldr_u(X9, X9, 0);	/* ldrb w9, [x9] */
		str_pre(X9, SP, -4);
		return;
	case 12: ldr_post(W0, SP, 4); E(0x13001c00); ldrx_u(X11, ST, 48); blr(X11); return; /* sxtb w0, w0 */
	}
	ldr_post(X9, SP, 4);
	ldr_u(X10, SP, 0);
	switch (k) {
	case 13: op3(0x0b000000, X10, X10, X9); break;
	case 14: op3(0x4b000000, X10, X10, X9); break;
	case 15: op3(0x1b007c00, X10, X10, X9); break;
	case 16: op3(0x0a000000, X10, X10, X9); break;
	case 17: op3(0x2a000000, X10, X10, X9); break;
	case 18: op3(0x4a000000, X10, X10, X9); break;
	case 19: op3(0x1ac02800, X10, X10, X9); break;
	default: E(0x6b09015f); E(0x1a9fa7ea);	/* cmp w10, w9; cset w10, lt */
	}
	str_u(X10, SP, 0);
}

static void *translate(int p)
{
	unsigned *start;
	int i = p, n = 0, w, k, t;
	if (cp - code > CODEW - 65536 || nbl >= MAXB) flush();
	start = cp;
	for (;;) {
		if (i >= MW) { E(0x12800000); b_to(x_exit); break; }	/* movn w0, #0: halt */
		if (n++ >= BLKW) { li(W0, i); b_to(x_disp); break; }
		w = m[i];
		cov[i] = 1;
		if (w & 1 && i + 1 < MW && !(m[i + 1] & 1) && (k = m[i + 1] / 2) >= 6 && k <= 8) {
			t = (w - 1) / 2;	/* literal + call/jmp/jz with a known target */
			cov[i + 1] = 1;
			if (k == 8) {
				unsigned *a;
				ldr_post(X10, SP, 4);
				a = cp; E(0x3500000a);	/* cbnz w10, fallthrough */
				li(W0, t); b_to(x_disp);
				cb_fix(a, cp);
				i += 2;
				continue;
			}
			li(X10, i + 2);
			if (k == 6) str_post(X10, RP, 4); else str_u(X10, RP, 0);
			li(W0, t); b_to(x_disp);
			i += 2;
			break;
		}
		if (w & 1) { li(X9, (w - 1) / 2); str_pre(X9, SP, -4); i++; continue; }
		k = w / 2;
		if (k == 0) { E(0x12800000); b_to(x_exit); i++; break; }
		if (k < 0 || k == 6 || k == 7) {
			ldr_post(W0, SP, 4);
			li(X10, i + 1);
			if (k == 6) str_post(X10, RP, 4); else str_u(X10, RP, 0);
			if (k < 0) { li(X11, 7 - k); addx_ext(RP, RP, X11, SXTW); }
			b_to(x_disp);
			i++;
			break;
		}
		if (k == 8) {	/* jz: taken leaves, not taken keeps going */
			ldp_post(X9, X10, SP, 8);
			E(0x3500006a);		/* cbnz w10, +3 */
			movw(W0, X9);
			b_to(x_disp);
			i++;
			continue;
		}
		if (k == 9) { ldr_pre(W0, RP, -4); b_to(x_disp); i++; break; }
		op(i, k);
		i++;
	}
	bl[nbl].a = p;
	bl[nbl++].b = i;
	tab[p] = start;
	syncicache(start, cp);
	if (trace) fprintf(stderr, "jit: %d..%d -> %ld bytes\n", p, i - 1, (long)((char *)cp - (char *)start));
	return start;
}

#else
/* portable path for hosts without the arm64 jit: a plain interpreter with
   the same semantics (shift counts wrap at 32, stores outside memory are
   ignored, loads outside memory read 0, leaving memory halts). */
#define LD(x) ((unsigned)(x) < MW ? m[x] : 0)
static void run(int p)
{
	int s = DS0, r = RS0, w, a, b;
	unsigned long long steps = 0;
	while ((unsigned)p < MW) {
		w = m[p++];
		steps++;
		if (w & 1) { m[--s] = w >> 1; continue; }
		a = m[s]; b = m[s + 1];
		switch (w / 2) {
		case 1: m[s] = LD(a); break;
		case 2: if ((unsigned)a < MW) m[a] = b; s += 2; break;
		case 3: m[s] = LD(s + 1 + a); break;
		case 4: s++; break;
		case 5: m[s] = b; m[s + 1] = a; break;
		case 6: m[r++] = p; p = a; s++; break;
		case 7: p = a; s++; break;
		case 8: if (!b) p = a; s += 2; break;
		case 9: p = m[--r]; break;
		case 10: m[--s] = (signed char)Q(); break;
		case 11: m[--s] = J; break;
		case 12: G((char)a); s++; break;
		case 13: m[++s] = b + a; break;
		case 14: m[++s] = (int)((unsigned)b - (unsigned)a); break;
		case 15: m[++s] = (int)((unsigned)b * (unsigned)a); break;
		case 16: m[++s] = b & a; break;
		case 17: m[++s] = b | a; break;
		case 18: m[++s] = b ^ a; break;
		case 19: m[++s] = b >> (a & 31); break;
		case 20: m[++s] = b < a; break;
		default: p = MW; break;	/* halt, and the reserved negative opcodes */
		}
	}
	if (trace) fprintf(stderr, "interp: %llu steps\n", steps);
}
#endif

int main(void)
{
	trace = getenv("ZYTRACE") != 0;
	memcpy(m + GORG, genome, sizeof genome);
#ifdef ZYJIT
	{
	int p = GORG;
	union { unsigned *u; int (*f)(long *, void *); } enter;
#ifdef __APPLE__
	code = mmap(0, CODEW * 4, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANON | MAP_JIT, -1, 0);
#else
	code = mmap(0, CODEW * 4, PROT_READ | PROT_WRITE | PROT_EXEC, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
#endif
	if (code == MAP_FAILED) { perror("mmap"); return 1; }
	jitwrite(1);
	build_stubs();
	jitwrite(0);
	syncicache(code, cp);
	S[0] = (long)m; S[1] = DS0; S[2] = RS0; S[3] = (long)tab; S[4] = (long)cov;
	S[5] = (long)Q; S[6] = (long)G; S[7] = (long)&J; S[8] = (long)hook;
	enter.u = x_enter;
	while ((unsigned)p < MW) {
		void *c = tab[p];
		if (!c) { jitwrite(1); c = translate(p); jitwrite(0); }
		p = enter.f(S, c);
	}
	}
#else
	run(GORG);
#endif
	return 0;
}