jitrie.c

/* 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;
}