compile.c

#include "moar.h"
#include "dasm_proto.h"
#include "platform/mmap.h"
#include "emit.h"

#define COPY_ARRAY(a, n, t) memcpy(MVM_malloc(n * sizeof(t)), a, n * sizeof(t))

static const MVMuint16 MAGIC_BYTECODE[] = { MVM_OP_sp_jit_enter, 0 };

MVMJitCode * MVM_jit_compile_graph(MVMThreadContext *tc, MVMJitGraph *jg) {
    dasm_State *state;
    char * memory;
    size_t codesize;
    /* Space for globals */
    MVMint32  num_globals = MVM_jit_num_globals();
    void ** dasm_globals = MVM_malloc(num_globals * sizeof(void*));
    MVMJitNode * node = jg->first_node;
    MVMJitCode * code;
    MVMint32 i;

    MVM_jit_log(tc, "Starting compilation\n");

    /* setup dasm */
    dasm_init(&state, 1);
    dasm_setupglobal(&state, dasm_globals, num_globals);
    dasm_setup(&state, MVM_jit_actions());
    dasm_growpc(&state, jg->labels_num);

    /* generate code */
    MVM_jit_emit_prologue(tc, jg,  &state);
    while (node) {
        switch(node->type) {
        case MVM_JIT_NODE_LABEL:
            MVM_jit_emit_label(tc, jg, &node->u.label, &state);
            break;
        case MVM_JIT_NODE_PRIMITIVE:
            MVM_jit_emit_primitive(tc, jg, &node->u.prim, &state);
            break;
        case MVM_JIT_NODE_BRANCH:
            MVM_jit_emit_branch(tc, jg, &node->u.branch, &state);
            break;
        case MVM_JIT_NODE_CALL_C:
            MVM_jit_emit_call_c(tc, jg, &node->u.call, &state);
            break;
        case MVM_JIT_NODE_GUARD:
            MVM_jit_emit_guard(tc, jg, &node->u.guard, &state);
            break;
        case MVM_JIT_NODE_INVOKE:
            MVM_jit_emit_invoke(tc, jg, &node->u.invoke, &state);
            break;
        case MVM_JIT_NODE_JUMPLIST:
            MVM_jit_emit_jumplist(tc, jg, &node->u.jumplist, &state);
            break;
        case MVM_JIT_NODE_CONTROL:
            MVM_jit_emit_control(tc, jg, &node->u.control, &state);
            break;
        }
        node = node->next;
    }
    MVM_jit_emit_epilogue(tc, jg, &state);

    /* compile the function */
    dasm_link(&state, &codesize);
    memory = MVM_platform_alloc_pages(codesize, MVM_PAGE_READ|MVM_PAGE_WRITE);
    dasm_encode(&state, memory);
    /* set memory readable + executable */
    MVM_platform_set_page_mode(memory, codesize, MVM_PAGE_READ|MVM_PAGE_EXEC);


    MVM_jit_log(tc, "Bytecode size: %"MVM_PRSz"\n", codesize);
    /* Create code segment */
    code = MVM_malloc(sizeof(MVMJitCode));
    code->func_ptr   = (MVMJitFunc)memory;
    code->size       = codesize;
    code->bytecode   = (MVMuint8*)MAGIC_BYTECODE;
    code->sf         = jg->sg->sf;

    /* Get the basic block labels */
    code->num_labels = jg->labels_num;
    code->labels = MVM_malloc(sizeof(void*) * code->num_labels);
    for (i = 0; i < code->num_labels; i++) {
        MVMint32 offset = dasm_getpclabel(&state, i);
        if (offset < 0)
            MVM_jit_log(tc, "Got negative offset for dynamic label %d\n", i);
        code->labels[i] = memory + offset;
    }

    /* Copy the deopts, inlines, and handlers. Because these use the label index
     * rather than the direct pointer, no fixup is necessary */
    code->num_bbs      = jg->bbs_num;
    code->bb_labels    = COPY_ARRAY(jg->bbs, jg->bbs_num, MVMint32);

    code->num_deopts   = jg->deopts_num;
    code->deopts       = code->num_deopts ? COPY_ARRAY(jg->deopts, jg->deopts_num, MVMJitDeopt) : NULL;
    code->num_handlers = jg->handlers_num;
    code->handlers     = code->num_handlers ? COPY_ARRAY(jg->handlers, jg->handlers_alloc, MVMJitHandler) : NULL;
    code->num_inlines  = jg->inlines_num;
    code->inlines      = code->num_inlines ? COPY_ARRAY(jg->inlines, jg->inlines_alloc, MVMJitInline) : NULL;

    /* clear up the assembler */
    dasm_free(&state);
    MVM_free(dasm_globals);

    if (tc->instance->jit_bytecode_dir) {
        MVM_jit_log_bytecode(tc, code);
    }
    if (tc->instance->jit_log_fh)
        fflush(tc->instance->jit_log_fh);
    return code;
}

void MVM_jit_destroy_code(MVMThreadContext *tc, MVMJitCode *code) {
    MVM_platform_free_pages(code->func_ptr, code->size);
    MVM_free(code);
}

#define X64_REGISTERS(_) \
    _(rax), \
    _(rcx), \
    _(rdx), \
    _(rbx), \
    _(rsp), \
    _(rbp), \
    _(rsi), \
    _(rdi), \
    _(r8), \
    _(r9), \
    _(r10), \
    _(r11), \
    _(r12), \
    _(r13), \
    _(r14), \
    _(r15)

typedef enum {
#define RENUM(x) x
X64_REGISTERS(RENUM)
#undef RENUM
} X64_REGISTER;

static const char * X64_REGISTER_NAMES[] = {
#define RNAME(x) #x
X64_REGISTERS(RNAME)
#undef RNAME
};


/* List of free registers, used as a 'register ring' */
static const X64_REGISTER FREE_REGISTERS[] = {
    rax, rcx, rdx, rsi, rdi, r8, r9, r10, r11
};

/* NB - valid only for POSIX call convention */
static const MVMint32 CALL_REGISTERS[] = {
    /* rdi, rsi, rdx, rcx, r8, r9 */
    4, 3, 2, 1, 5, 6
};

#define NUM_REGISTERS (sizeof(FREE_REGISTERS)/sizeof(X64_REGISTER))
#define NUM_CALL_REGS (sizeof(FREE_REGISTERS)/sizeof(X64_REGISTER))

typedef struct {
    MVMint32 reg_used[NUM_REGISTERS]; /* last use of a register */
    MVMint32 *nodes_reg;      /* register number of node */
    MVMint32 *nodes_spill;    /* spill location of a given node, if spilled */
    MVMint32 last_register;   /* last register number allocated; used
                                 to implement the 'register ring' */
    MVMint32 last_spill;      /* last spill threshold, used to ensure
                               * that every value receives it's own
                               * node */
    MVMint32 cond_depth;      /* last used local label depth by
                                 conditional statement */
} CompilerRegisterState;

static void take_register(MVMThreadContext *tc, CompilerRegisterState *state, MVMint32 regnum) {
    MVMint32 node, spill;
    if (state->reg_used[regnum] >= 0) {
        /* register is used. spill it's value */
        node  = state->reg_used[regnum];
        spill = state->nodes_spill[node];

        if (spill < 0) {
            /* node was never spilled before, allocate a location  */
            spill = state->last_spill++;
        }
        MVM_jit_log(tc, "mov [rsp+0x%x], %s\n", spill*sizeof(MVMRegister),
                    X64_REGISTER_NAMES[FREE_REGISTERS[regnum]]);
        /* mark node as spilled */
        state->nodes_reg[node]        = -1;
        state->nodes_spill[node]      = spill;
    }
    /* mark register as free */
    state->reg_used[regnum] = -1;
}


static void emit_full_spill(MVMThreadContext *tc, CompilerRegisterState *state) {
    /* Equivalent to simply taking all registers */
    MVMint32 i;
    for (i = 0; i < NUM_REGISTERS; i++) {
        take_register(tc, state, i);
    }
}

/* invalidate registers */
static void invalidate_registers(MVMThreadContext *tc, CompilerRegisterState *state) {
    MVMint32 i;
    for (i = 0; i < NUM_REGISTERS; i++) {
        MVMint32 node = state->reg_used[i];
        if (node >= 0) {
            state->nodes_reg[node] = -1;
            state->reg_used[i]     = -1;
        }
    }
}

static void load_node_to(MVMThreadContext *tc, CompilerRegisterState *state, MVMint32 node, MVMint32 regnum) {
    MVMint32 cur_reg = state->nodes_reg[node];
    if (state->reg_used[regnum] == node)
        return;
    if (cur_reg >= 0) {
        MVM_jit_log(tc, "mov %s, %s\n", X64_REGISTER_NAMES[FREE_REGISTERS[regnum]],
                    X64_REGISTER_NAMES[FREE_REGISTERS[cur_reg]]);
        state->reg_used[cur_reg] = -1; /* is now free */
    } else if (state->nodes_spill[node] >= 0) {
        MVM_jit_log(tc, "mov %s, [rsp+0x%x]\n", X64_REGISTER_NAMES[FREE_REGISTERS[regnum]],
                    state->nodes_spill[node]*sizeof(MVMRegister));
    } else {
        MVM_oops(tc, "Requested load of node %d but not in register or memory\n");
    }
    state->nodes_reg[node]  = regnum;
    state->reg_used[regnum] = node;
}

static MVMint32 get_next_register(MVMThreadContext *tc, CompilerRegisterState *state,
                                  X64_REGISTER *regs, MVMint32 nregs) {
    MVMint32 i, j;
    for (i = 0; i < NUM_REGISTERS; i++) {
        MVMint32 regnum = (state->last_register++) % NUM_REGISTERS;
        for (j = 0; j < nregs; j++) {
            if (regs[j] == FREE_REGISTERS[regnum])
                break;
        }
        if (j < nregs)
            continue;
        /* register was not found in input. if necessary, steal it */
        take_register(tc, state, regnum);
        return regnum;
    }
    MVM_oops(tc, "Could not allocate a register\n");
}

static void emit_expr_op(MVMThreadContext *tc, CompilerRegisterState *state, MVMJitExprNode op,
                         X64_REGISTER *regs, MVMJitExprNode *args) {
    switch(op) {
    case MVM_JIT_LOAD:
        MVM_jit_log(tc, "mov %s, [%s]\n", X64_REGISTER_NAMES[regs[1]], X64_REGISTER_NAMES[regs[0]]);
        break;
    case MVM_JIT_STORE:
        MVM_jit_log(tc, "mov [%s], %s\n", X64_REGISTER_NAMES[regs[0]], X64_REGISTER_NAMES[regs[1]]);
        break;
    case MVM_JIT_COPY:
        MVM_jit_log(tc, "mov %s, %s\n", X64_REGISTER_NAMES[regs[1]], X64_REGISTER_NAMES[regs[0]]);
        break;
    case MVM_JIT_ADDR:
        MVM_jit_log(tc, "lea %s, [%s+0x%"PRIx64"]\n", X64_REGISTER_NAMES[regs[1]],
                    X64_REGISTER_NAMES[regs[0]], args[0]);
        break;
    case MVM_JIT_IDX:
        MVM_jit_log(tc, "lea %s, [%s+%s*%d]\n", X64_REGISTER_NAMES[regs[2]],
                    X64_REGISTER_NAMES[regs[0]], X64_REGISTER_NAMES[regs[1]], args[0]);
        break;
    case MVM_JIT_ADD:
        MVM_jit_log(tc, "mov %s, %s\n", X64_REGISTER_NAMES[regs[2]], X64_REGISTER_NAMES[regs[0]]);
        MVM_jit_log(tc, "add %s, %s\n", X64_REGISTER_NAMES[regs[2]], X64_REGISTER_NAMES[regs[1]]);
        break;
    case MVM_JIT_AND:
        MVM_jit_log(tc, "mov %s, %s\n", X64_REGISTER_NAMES[regs[2]], X64_REGISTER_NAMES[regs[0]]);
        MVM_jit_log(tc, "and %s, %s\n", X64_REGISTER_NAMES[regs[2]], X64_REGISTER_NAMES[regs[1]]);
        break;
    case MVM_JIT_NZ:
        take_register(tc, state, 0);
        MVM_jit_log(tc, "test %s, %s\n", X64_REGISTER_NAMES[regs[0]], X64_REGISTER_NAMES[regs[0]]);
        MVM_jit_log(tc, "setnz al\n");
        MVM_jit_log(tc, "movzx %s, al\n", X64_REGISTER_NAMES[regs[1]]);
        break;
    case MVM_JIT_ZR:
        take_register(tc, state, 0);
        MVM_jit_log(tc, "test %s, %s\n", X64_REGISTER_NAMES[regs[0]], X64_REGISTER_NAMES[regs[0]]);
        MVM_jit_log(tc, "setnz al\n");
        MVM_jit_log(tc, "movzx %s, al\n", X64_REGISTER_NAMES[regs[1]]);
        break;
    case MVM_JIT_LOCAL:
        MVM_jit_log(tc, "mov %s, %s\n", X64_REGISTER_NAMES[regs[0]], X64_REGISTER_NAMES[rbx]);
        break;
    case MVM_JIT_FRAME:
        MVM_jit_log(tc, "mov %s, %s\n", X64_REGISTER_NAMES[regs[0]], X64_REGISTER_NAMES[r12]);
        break;
    case MVM_JIT_CU:
        MVM_jit_log(tc, "mov %s, %s\n", X64_REGISTER_NAMES[regs[0]], X64_REGISTER_NAMES[r13]);
        break;
    case MVM_JIT_TC:
        MVM_jit_log(tc, "mov %s, %s\n", X64_REGISTER_NAMES[regs[0]], X64_REGISTER_NAMES[r14]);
        break;
    case MVM_JIT_CONST:
        MVM_jit_log(tc, "mov %s, 0x%x\n", X64_REGISTER_NAMES[regs[0]], (MVMint32)args[0]);
        break;
    case MVM_JIT_CALL:
        MVM_jit_log(tc, "call %s\n", X64_REGISTER_NAMES[regs[0]]);
        break;
    default: {
        const MVMJitExprOpInfo *info = MVM_jit_expr_op_info(tc, op);
        MVM_jit_log(tc, "not yet sure how to compile %s\n", info->name);
    }
    }
}

static void load_op_regs(MVMThreadContext *tc, CompilerRegisterState *state,
                         MVMJitExprTree *tree, MVMint32 node, X64_REGISTER *regs) {
    MVMJitExprNode op            = tree->nodes[node];
    const MVMJitExprOpInfo *info = MVM_jit_expr_op_info(tc, op);
    MVMint32 i, first_child, nchild;
    first_child = node + 1;
    nchild      = (info->nchild < 0 ? tree->nodes[first_child++] : info->nchild);
    /* ensure child nodes have been computed into memory */
    for (i = 0; i < nchild; i++) {
        MVMint32 child  = tree->nodes[first_child+i];
        MVMint32 regnum = state->nodes_reg[child];
        if (MVM_jit_expr_op_info(tc, tree->nodes[child])->vtype == MVM_JIT_VOID)
            continue;
        if (regnum < 0) {
            regnum = get_next_register(tc, state, regs, i);
            load_node_to(tc, state, child, regnum);
        }
        regs[i] = FREE_REGISTERS[regnum];
    }
    if (info->vtype != MVM_JIT_VOID) {
        /* assign an output register */
        MVMint32 regnum = get_next_register(tc, state, regs, i);
        regs[i]         = FREE_REGISTERS[regnum];
        state->nodes_reg[node]  = regnum;
        state->reg_used[regnum] = node;
    }
}

static void load_call_regs(MVMThreadContext *tc, CompilerRegisterState *state,
                           MVMJitExprTree *tree, MVMint32 node, X64_REGISTER *regs) {
    MVMint32 arglist = tree->nodes[node+2];
    MVMint32 nargs   = tree->nodes[arglist+1];
    MVMint32 i, func;
    for (i = 0; i < nargs; i++) {
        MVMint32 carg   = tree->nodes[arglist+2+i];
        MVMint32 argval = tree->nodes[carg+1];
        MVMint32 argtyp = tree->nodes[carg+2];
        MVMint32     regnum  = CALL_REGISTERS[i];
        X64_REGISTER reg     = FREE_REGISTERS[regnum];
        /* whatever, ignore argtyp */
        load_node_to(tc, state, argval, regnum);
        regs[i] = reg;
    }
    /* load function register */
    func = get_next_register(tc, state, regs, nargs);
    load_node_to(tc, state, tree->nodes[node+1], func);
    regs[0] = FREE_REGISTERS[func];
}

static void prepare_expr_op(MVMThreadContext *tc, MVMJitTreeTraverser *traverser,
                            MVMJitExprTree *tree, MVMint32 node) {
    /* Spill before call or if or when */
    CompilerRegisterState *state = traverser->data;
    MVMJitExprNode op = tree->nodes[node];
    if (traverser->visits > 0)
        return;
    /* Conditional blocks should spill before they are run */
    switch (op) {
    case MVM_JIT_WHEN:
        /* require a label for the statement end */
        state->cond_depth += 1;
        emit_full_spill(tc, state);
        break;
    case MVM_JIT_IF:
        /* require two labels, one for the alternative block, and one for the statement end */
        state->cond_depth += 2;
        emit_full_spill(tc, state);
        break;
    default:
        break;
    }
}

static void compile_expr_labels(MVMThreadContext *tc, MVMJitTreeTraverser *traverser,
                                MVMJitExprTree *tree, MVMint32 node, MVMint32 i) {
    CompilerRegisterState *state = traverser->data;
    MVMJitExprNode op = tree->nodes[node];
    if (traverser->visits > 0)
        return;
    switch (op) {
    case MVM_JIT_IF:
        /* 'ternary operator' or 'expression style' if */
        if (i == 0) {
            /* branch to second option (label 1) */
            MVM_jit_log(tc, "jnz >%d\n", state->cond_depth - 1);
        } else if (i == 1) {
            /* move result value into place (i.e. rax) */
            load_node_to(tc, state, node + 2, 0);
            /* just after first option, branch to end */
            MVM_jit_log(tc, "jmp >%d\n", state->cond_depth);
            MVM_jit_log(tc, "%d:\n" , state->cond_depth - 1);
            /* registers used in the 'then' block are not usable in the 'else' block */
            invalidate_registers(tc, state);
        } else {
            /* move result value into place */
            load_node_to(tc, state, node + 3, 0);
            /* emit end label */
            MVM_jit_log(tc, "%d:\n", state->cond_depth);
            invalidate_registers(tc, state);
        }
        break;
    case MVM_JIT_WHEN:
        if (i == 0) {
            MVM_jit_log(tc, "jnz >%d\n", state->cond_depth);
        } else {
            MVM_jit_log(tc, "%d:\n", state->cond_depth);
            invalidate_registers(tc, state);
        }
        break;
    }
}

static void compile_expr_op(MVMThreadContext *tc, MVMJitTreeTraverser *traverser,
                            MVMJitExprTree *tree, MVMint32 node) {
    CompilerRegisterState *state = traverser->data;
    MVMJitExprNode op  = tree->nodes[node];
    const MVMJitExprOpInfo *info = MVM_jit_expr_op_info(tc, op);
    X64_REGISTER regs[8];
    MVMJitExprNode *args = tree->nodes + node + (info->nchild < 0 ? tree->nodes[node+1] + 1 : info->nchild) + 1;
    if (traverser->visits[node] > 1) /* no revisits */
        return;
    switch(op) {
    case MVM_JIT_CALL:
        /* spill before we go */
        emit_full_spill(tc, state);
        load_call_regs(tc, state, tree, node, regs);
        emit_expr_op(tc, state, op, regs, args);
        /* all loaded registers are now invalid */
        invalidate_registers(tc, state);
        if (args[0] != MVM_JIT_RV_VOID) {
            state->nodes_reg[node]  = 0;
            state->reg_used[0]      = node;
        }
        break;
    case MVM_JIT_WHEN:
        state->cond_depth     -= 1;
        break;
    case MVM_JIT_IF:
        /* result value lives in rax (as assured in compile_expr_labels) */
        state->reg_used[0]     = node;
        state->nodes_reg[node] = 0;
        state->cond_depth     -= 2;
        break;
    case MVM_JIT_DO:
        {
            MVMint32 nchild     = tree->nodes[node+1];
            MVMint32 last_child = tree->nodes[node+1+nchild];
            MVMint32 regnum     = get_next_register(tc, state, regs, 0);
            load_node_to(tc, state, last_child, regnum);
            state->reg_used[regnum] = node;
            state->nodes_reg[node]  = regnum;
            break;
        }
    case MVM_JIT_ALL:
        {
            MVMint32 nchild      = tree->nodes[node+1];
            MVMint32 first_child = tree->nodes[node+2];
            MVMint32 result, value, i;
            if (nchild == 0)
                MVM_oops(tc, "No child for ALL is clearly an error");
            result  = get_next_register(tc, state, regs, 0);
            regs[0] = FREE_REGISTERS[result];
            value   = get_next_register(tc, state, regs, 1);
            /* load first child */
            load_node_to(tc, state, first_child, result);
            for (i = 1; i < nchild; i++) {
                MVMint32 child = tree->nodes[node+i+2];
                load_node_to(tc, state, child, value);
                MVM_jit_log(tc, "and %s, %s\n", X64_REGISTER_NAMES[FREE_REGISTERS[result]],
                            X64_REGISTER_NAMES[FREE_REGISTERS[value]]);
            }
            state->nodes_reg[node]  = result;
            state->reg_used[result]  = node;
            break;
        }
    case MVM_JIT_CARG:
    case MVM_JIT_ARGLIST:
        break;
    default:
        load_op_regs(tc, state, tree, node, regs);
        emit_expr_op(tc, state, op, regs, args);
        break;
    }

}


void MVM_jit_compile_expr_tree(MVMThreadContext *tc, MVMJitGraph *jg, MVMJitExprTree *tree) {
    /* VERY SIMPLE AND PROVISIONAL COMPILER FOR THE EXPRESSION TREE
     *
     * We use the tree nodes as value numbers, and compile everything
     * to registers. Spill to stack (much simpler that way, but wrong
     * with regard to GC). This algorithm is a proof of concept, but
     * the concept is to do online register allocation without any use
     * knowledge, which is not really a good concept at all. */
    MVMJitTreeTraverser compiler;
    CompilerRegisterState state;
    /* Initialize state */
    memset(state.reg_used, -1, sizeof(state.reg_used));
    state.nodes_reg     = MVM_malloc(sizeof(MVMint32)*tree->nodes_num);
    state.nodes_spill   = MVM_malloc(sizeof(MVMint32)*tree->nodes_num);
    state.last_register = 0;
    state.last_spill    = 0;
    state.cond_depth    = 0;
    memset(state.nodes_reg, -1, sizeof(MVMint32)*tree->nodes_num);
    memset(state.nodes_spill, -1, sizeof(MVMint32)*tree->nodes_num);

    /* initialize compiler */
    memset(&compiler, 0, sizeof(MVMJitTreeTraverser));
    compiler.data      = &state;
    compiler.preorder  = &prepare_expr_op;
    compiler.inorder   = &compile_expr_labels;
    compiler.postorder = &compile_expr_op;

    MVM_jit_expr_tree_traverse(tc, tree, &compiler);

    MVM_free(state.nodes_reg);
    MVM_free(state.nodes_spill);
}


/* Returns 1 if we should return from the frame, the function, 0 otherwise */
MVMint32 MVM_jit_enter_code(MVMThreadContext *tc, MVMCompUnit *cu,
                            MVMJitCode *code) {
    /* The actual JIT code returns 0 if it went through to the exit */
    void *label = tc->cur_frame->jit_entry_label;
    MVMint32 ctrl = code->func_ptr(tc, cu, label);
    return ctrl ? 0 : 1;
}