#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>

enum instr {
	PushRef,
	PushI,
	Put,
	Pop,

	Call,
	Jmp,
	JmpTrue,
	Ret,
	Halt,

	IAddRet,
	IMulRet,
	ISubRet,
	IDivRet
};

static inline uint32_t instr_size(enum instr instr) {
	switch (instr) {
		case PushRef: return 5+1;
		case PushI:   return 7+1;
		case Put:     return 1+5;
		case Pop:     return 4;

		case Call:    return 5;
		case Jmp:     return 5;
		case JmpTrue: return 1+3+6;
		case Ret:     return 1;
		case Halt:    return 1+1;

		case IAddRet:
		case ISubRet: return 5+5+3+5+1;
		case IMulRet: return 5+5+4+5+1;
		case IDivRet: return 5+5+3+3+5+1;

		default:
			fprintf(stderr,"unknown instruction %d\n",instr);
			exit(1);
	}
}

static inline void gen_instr(char *full_code, char **code_p, uint64_t **pgm_p, uint32_t *mapping) {
	uint64_t arg;

	char *code=*code_p;
	uint64_t *pgm=*pgm_p;

	switch (*pgm) {
		case PushRef:
			arg=pgm[1];
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"PushRef %ld\n",arg);
#endif
			code[0]='\x48'; /* mov rcx,[rsp+ARG*8] */
			code[1]='\x8b';
			code[2]='\x4c';
			code[3]='\x24';
			code[4]=(unsigned char)arg*8;
			code[5]='\x51'; /* push rcx */
			pgm+=2;
			code+=6;
			break;
		case PushI:
			arg=pgm[1];
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"PushI %lu\n",arg);
#endif
			code[0]='\x48'; /* mov rcx,ARG */
			code[1]='\xc7';
			code[2]='\xc1';
			*(uint32_t*)&code[3]=(uint32_t)arg;
			code[7]='\x51'; /* push rcx */
			pgm+=2;
			code+=8;
			break;
		case Put:
			arg=pgm[1];
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"Put %lu\n",arg);
#endif
			code[0]='\x59'; /* pop rcx */
			code[1]='\x48'; /* mov [rsp+ARG*8],rcx */
			code[2]='\x89';
			code[3]='\x4c';
			code[4]='\x24';
			code[5]=(unsigned char)arg*8;
			pgm+=2;
			code+=6;
			break;
		case Pop:
			arg=pgm[1];
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"Pop %lu\n",arg);
#endif
			code[0]='\x48'; /* add rsp,ARG*8 */
			code[1]='\x83';
			code[2]='\xc4';
			code[3]=(unsigned char)arg*8;
			pgm+=2;
			code+=4;
			break;

		case Call:
			arg=pgm[1];
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"Call %lu -> %d\n",arg,mapping[arg]-(uint32_t)(&code[5]-full_code));
#endif
			code[0]='\xe8'; /* callq ARG */
			*(uint32_t*)&code[1]=mapping[arg]-(&code[5]-full_code);
			pgm+=2;
			code+=5;
			break;
		case Jmp:
			arg=pgm[1];
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"Jmp %lu -> %d\n",arg,mapping[arg]-(uint32_t)(&code[5]-full_code));
#endif
			code[0]='\xe9'; /* jmpq ARG */
			*(uint32_t*)&code[1]=mapping[arg]-(&code[5]-full_code);
			pgm+=2;
			code+=5;
			break;
		case JmpTrue:
			arg=pgm[1];
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"JmpTrue %lu -> %d\n",arg,mapping[arg]-(uint32_t)(&code[10]-full_code));
#endif
			code[0]='\x59'; /* pop rcx */
			code[1]='\x48'; /* test rcx,rcx */
			code[2]='\x85';
			code[3]='\xc9';
			code[4]='\x0f'; /* jne ARG */
			code[5]='\x85';
			*(uint32_t*)&code[6]=mapping[arg]-(&code[10]-full_code);
			pgm+=2;
			code+=10;
			break;
		case Ret:
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"Ret\n");
#endif
			code[0]='\xc3'; /* retq */
			pgm++;
			code+=1;
			break;
		case Halt:
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"Halt\n");
#endif
			code[0]='\x58'; /* pop rax */
			code[1]='\xc3'; /* retq */
			pgm++;
			code+=2;
			break;

		case IAddRet:
		case IMulRet:
		case ISubRet:
		case IDivRet:
#ifdef DEBUG_JIT_INSTRUCTIONS
			fprintf(stderr,"I<Op>Ret\n");
#endif
			/* mov rax,[rsp+8] */
			code[0]='\x48'; code[1]='\x8b'; code[2]='\x44'; code[3]='\x24'; code[4]='\x08';
			/* mov rcx,[rsp+16] */
			code[5]='\x48'; code[6]='\x8b'; code[7]='\x4c'; code[8]='\x24'; code[9]='\x10';
			switch (*pgm) {
				case IAddRet:
				case ISubRet:
					/* {add,sub} rax,rcx */
					code[10]='\x48';
					code[11]=*pgm==IAddRet ? '\x01' : '\x29';
					code[12]='\xc8';
					code+=13;
					break;
				case IMulRet:
					/* imul rax,rcx */
					code[10]='\x48'; code[11]='\x0f'; code[12]='\xaf'; code[13]='\xc1';
					code+=14;
					break;
				case IDivRet:
					/* xor rdx,rdx */
					code[10]='\x48'; code[11]='\x31'; code[12]='\xd2';
					/* idiv rcx */
					code[13]='\x48'; code[14]='\xf7'; code[15]='\xf9';
					code+=16;
					break;
			}
			/* mov [rsp+16],rax */
			code[0]='\x48'; code[1]='\x89'; code[2]='\x44'; code[3]='\x24'; code[4]='\x10';
			/* retq */
			code[5]='\xc3';
			pgm++;
			code+=6;
			break;

		default:
			fprintf(stderr,"unknown instruction %d\n",(int)pgm[-1]);
			exit(1);
	}

	*code_p=code;
	*pgm_p=pgm;
}

void init_jit(uint64_t max_instrs, uint64_t max_code_len,
		char **code_block, uint32_t **mapping) {
	*code_block=mmap (NULL,max_code_len,PROT_READ | PROT_EXEC,MAP_ANONYMOUS | MAP_PRIVATE,0,0);
	*mapping=malloc (max_instrs*sizeof(uint32_t));
}

char *jit_append(char *code_block, uint32_t code_len, char *code_ptr,
		uint32_t *mapping, uint32_t n_instr, uint64_t *pgm, int is_main) {
	uint32_t len=*pgm++;
	uint32_t i;
	uint32_t code_i=code_ptr-code_block;

	uint64_t *pgm_p=pgm;
	for (i=n_instr; i<n_instr+len; i++) {
		enum instr instr=(enum instr)*pgm_p;

		mapping[i]=code_i;
		code_i+=instr_size(instr);

		switch (instr) {
			case PushRef:
			case PushI:
			case Put:
			case Pop:
			case Call:
			case Jmp:
			case JmpTrue:
				pgm_p+=2;
				break;
			default:
				pgm_p+=1;
				break;
		}
	}

	mprotect (code_block,code_len,PROT_WRITE);

	if (is_main)
		*(uint32_t*)&code_block[9]=code_ptr-code_block-13;

	pgm_p=pgm;
	for (i=0; i<len; i++)
		gen_instr (code_block,&code_ptr,&pgm_p,mapping);

	mprotect (code_block,code_len,PROT_READ | PROT_EXEC);

	return code_ptr;
}

uint64_t jit_exec(char *code_block) {
	return ((uint64_t(*)(void))code_block)();
}