Create a gist now

Instantly share code, notes, and snippets.

@lpereira /jitrie.c Secret
Created Apr 13, 2014

What would you like to do?
/* First jitrie implementation.
* Leandro A. F. Pereira, April 2014
*
* This isn't really creating code to traverse a tree. Also, the generated machine code isn't
* still working (need to fix the relative jump offsets). But the C code works.
*/
#include <stdint.h>
#include <string.h>
#include <stdio.h>
#include <sys/mman.h>
struct entry {
const char *str;
int value;
} test_entries[] = {
{ "/hello", 0 },
{ "/chunked", 1 },
{ "/sse", 2 },
{ "/beacon", 3 },
{ "/favicon.ico", 4 },
{ "/admin", 5 },
{ "/", 6 },
};
char *opcode_to_str[] = {
"NOP",
"PROLOG",
"EPILOG",
"STR_CMP_4",
"STR_CMP_2",
"STR_CMP_1",
"INC_OFFSET",
"RETURN",
};
typedef enum {
NOP,
PROLOG,
EPILOG,
STR_CMP_4,
STR_CMP_2,
STR_CMP_1,
INC_OFFSET,
RETURN,
} opcode;
typedef struct {
opcode op;
int32_t val;
} insn;
insn *emit(insn *ir, opcode op, int32_t val)
{
ir->op = op;
ir->val = val;
return ir + 1;
}
insn *jit_prolog(insn *ir)
{
ir = emit(ir, PROLOG, 0);
return ir;
}
insn *jit_epilog(insn *ir)
{
ir = emit(ir, RETURN, -1);
ir = emit(ir, EPILOG, 0);
return ir;
}
int32_t string4_to_int32(const char *s)
{
return (int32_t)(s[0] | s[1] << 8 | s[2] << 16 | s[3] << 24);
}
int16_t string2_to_int16(const char *s)
{
return (int16_t)(s[0] | s[1] << 8);
}
insn *jit_entry(insn *ir, struct entry *entry)
{
const char *str = entry->str;
size_t current_len = strlen(entry->str);
while (current_len != 0) {
if (current_len >= 4) {
ir = emit(ir, STR_CMP_4, string4_to_int32(str));
ir = emit(ir, INC_OFFSET, 4);
str += 4;
current_len -= 4;
} else if (current_len >= 2) {
ir = emit(ir, STR_CMP_2, string2_to_int16(str));
ir = emit(ir, INC_OFFSET, 2);
str += 2;
current_len -= 2;
} else {
ir = emit(ir, STR_CMP_1, *str);
ir = emit(ir, INC_OFFSET, 1);
str++;
current_len--;
}
if (!*str) {
ir = emit(ir, RETURN, entry->value);
ir = emit(ir, INC_OFFSET, -strlen(entry->str));
}
}
return ir;
}
insn *jit(insn *ir, struct entry *entries, size_t size)
{
size_t i;
ir = jit_prolog(ir);
for (i = 0; i < size; i++) {
ir = jit_entry(ir, &entries[i]);
}
ir = jit_epilog(ir);
return ir;
}
void gen_c(insn *ir_start, insn *ir_end)
{
insn *i;
int offset = 0;
for (i = ir_start; i < ir_end; i++) {
switch (i->op) {
case NOP:
printf("/* NOP */\n");
break;
case PROLOG:
printf("#include <stdint.h>\n");
printf("int j(char *key) {\n");
break;
case EPILOG:
printf("}\n");
break;
case STR_CMP_4:
printf("if (*((int32_t *) (key + %d)) == %d) ", offset, i->val);
break;
case STR_CMP_2:
printf("if (*((int16_t *) (key + %d)) == %d) ", offset, i->val);
break;
case STR_CMP_1:
printf("if (*(key + %d) == '%c') ", offset, i->val);
break;
case INC_OFFSET:
offset += i->val;
break;
case RETURN:
printf("return %d;\n", i->val);
break;
}
}
}
void *gen_x86_64_asm(insn *ir_start, insn *ir_end, size_t *s)
{
/* The generated code here is far from being optimal, but should do as
* an a proof of concept. */
insn *i;
int offset = 0;
unsigned char *ptr = mmap(0, 4096, PROT_READ | PROT_WRITE | PROT_EXEC,
MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
unsigned char *ptr_orig = ptr;
unsigned char *last_jne_loc = NULL;
unsigned char *ret_block_locs[64];
int ret_block_locs_idx = 0;
if (ptr == MAP_FAILED) {
perror("mmap");
return NULL;
}
*s = 0;
#define EMIT(byte) do { *ptr++ = (byte); *s = *s + 1; } while(0)
for (i = ir_start; i < ir_end; i++) {
switch (i->op) {
case NOP:
break;
case PROLOG:
// push %rbp
EMIT(0x55);
// mov %rsp, %rbp
EMIT(0x48);
EMIT(0x89);
EMIT(0xe5);
// mov %rdi, -0x8(%rbp)
EMIT(0x48);
EMIT(0x89);
EMIT(0x7d);
EMIT(0xf8);
break;
case EPILOG:
for (ret_block_locs_idx--; ret_block_locs_idx; ret_block_locs_idx--) {
*ret_block_locs[ret_block_locs_idx] = (unsigned char)(ptr - ret_block_locs[ret_block_locs_idx]);
}
// pop %rbp
EMIT(0x5d);
// retq
EMIT(0xc3);
break;
case STR_CMP_4:
// mov -0x8(%rbp), %rax
EMIT(0x48);
EMIT(0x8b);
EMIT(0x45);
EMIT(0xf8);
if (offset) {
// add $offset, $rax
EMIT(0x48);
EMIT(0x83);
EMIT(0xc0);
EMIT(offset & 0xff);
}
// mov (%rax), %eax
EMIT(0x8b);
EMIT(0x00);
// cmp val, %eax
EMIT(0x3d);
EMIT(i->val & 0xff);
EMIT(i->val >> 8 & 0xff);
EMIT(i->val >> 16 & 0xff);
EMIT(i->val >> 24 & 0xff);
// jne next
EMIT(0x75);
last_jne_loc = ptr;
EMIT(0x00);
break;
case STR_CMP_2:
// mov -0x8(%%rbp), %%rax
EMIT(0x48);
EMIT(0x8b);
EMIT(0x45);
EMIT(0xf8);
if (offset) {
// add $0x%x,%%rax
EMIT(0x48);
EMIT(0x83);
EMIT(0xc0);
EMIT(offset & 0xff);
}
// movzwl (%%rax),%%eax
EMIT(0x0f);
EMIT(0xb7);
EMIT(0x00);
// cmp $0x%x, %%eax
EMIT(0x66);
EMIT(0x3d);
EMIT(i->val & 0xff);
EMIT(i->val >> 8 & 0xff);
// jne next
EMIT(0x75);
last_jne_loc = ptr;
EMIT(0x00);
break;
case STR_CMP_1:
// mov -0x8(%%rbp), %%rax
EMIT(0x48);
EMIT(0x8b);
EMIT(0x45);
EMIT(0xf8);
if (offset) {
// add $0x%x,%%rax
EMIT(0x48);
EMIT(0x83);
EMIT(0xc0);
EMIT(offset & 0xff);
}
// movzbl (%%rax),%%eax
EMIT(0x0f);
EMIT(0xb6);
EMIT(0x00);
// cmp val, %%eax
EMIT(0x3c);
EMIT(i->val & 0xff);
// jne next
EMIT(0x75);
last_jne_loc = ptr;
EMIT(0x00);
break;
case INC_OFFSET:
offset += i->val;
if (offset == 0) {
*last_jne_loc = 0x42; //(unsigned char)(ptr - last_jne_loc);
}
break;
case RETURN:
// mov val, %eax
EMIT(0xb8);
EMIT(i->val & 0xff);
EMIT(i->val >> 8 & 0xff);
EMIT(i->val >> 16 & 0xff);
EMIT(i->val >> 24 & 0xff);
// jmp end
EMIT(0xeb);
EMIT(0x00);
ret_block_locs[ret_block_locs_idx] = ptr;
ret_block_locs_idx++;
break;
}
}
return ptr_orig;
}
int main(void)
{
insn ir_alloc[2048];
insn *ir;
ir = jit(ir_alloc, test_entries, sizeof(test_entries) / sizeof(test_entries[0]));
printf("/* C code */\n");
gen_c(ir_alloc, ir);
printf("/* x86_64 asm code */\n");
size_t s;
int (*jitted_fn)(const char *k) = gen_x86_64_asm(ir_alloc, ir, &s);
if (!jitted_fn)
return 1;
printf("*** code size is %ld\n", s);
size_t i;
unsigned char *bla = (unsigned char *)jitted_fn;
for (i = 0; i < s; i++) {
printf("%02x ", bla[i]);
if (i % 16 == 0 && i != 0)
putchar('\n');
}
return 0;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment