Stub logic working for fib test, but still crashing in other cases

3 files changed, 71 insertions, 39 deletions

diff --git a/ujit_codegen.c b/ujit_codegen.c
index e39e418979..215bc3cf92 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -69,15 +69,6 @@ ujit_side_exit(codeblock_t* cb, ctx_t* ctx)
     // Table mapping opcodes to interpreter handlers
     const void * const *table = rb_vm_get_insns_address_table();
 
-    // Write back the old instruction at the entry PC
-    // To deotimize the code block this instruction belongs to
-    VALUE* entry_pc = &ctx->iseq->body->iseq_encoded[ctx->start_idx];
-    int entry_opcode = opcode_at_pc(ctx->iseq, entry_pc);
-    void* entry_instr = (void*)table[entry_opcode];
-    mov(cb, RAX, const_ptr_opnd(entry_pc));
-    mov(cb, RCX, const_ptr_opnd(entry_instr));
-    mov(cb, mem_opnd(64, RAX, 0), RCX);
-
     // Write back the old instruction at the exit PC
     // Otherwise the interpreter may jump right back to the
     // JITted code we're trying to exit
@@ -99,7 +90,7 @@ Compile a sequence of bytecode instructions starting at `insn_idx`.
 Returns `NULL` if compilation fails.
 */
 uint8_t *
-ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, bool gen_entry)
+ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, bool entry_point)
 {
     assert (cb != NULL);
     VALUE *encoded = iseq->body->iseq_encoded;
@@ -148,9 +139,9 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, bool gen_entry)
             break;
         }
 
-        // If requested, write the interpreter entry
-        // prologue before the first instruction
-        if (gen_entry && num_instrs == 0) {
+        // If this is an interpreter entry point, write the interpreter
+        // entry prologue before the first instruction
+        if (entry_point && num_instrs == 0) {
             ujit_gen_entry(cb);
         }
 
@@ -172,15 +163,20 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, bool gen_entry)
         }
     }
 
+    // FIXME: only generate exit if no instructions were compiled?
+    // or simply don't allow instructions to fail to compile anymore?
     // If no instructions were compiled
-    if (num_instrs == 0) {
-        return NULL;
-    }
+    //if (num_instrs == 0) {
+    //    return NULL;
+    //}
 
     // Generate code to exit to the interpreter
     ujit_gen_exit(cb, &ctx, &encoded[insn_idx]);
 
-    map_addr2insn(code_ptr, first_opcode);
+    // If this is an interpreter entry point
+    if (entry_point) {
+        map_addr2insn(code_ptr, first_opcode);
+    }
 
     if (UJIT_DUMP_MODE >= 2) {
         // Dump list of compiled instrutions
@@ -897,13 +893,15 @@ static bool
 gen_branchunless(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx)
 {
     // Get the branch target instruction offsets
-    int32_t jump_idx = (int32_t)ctx_get_arg(ctx, 0);
-    int32_t next_idx = ctx->insn_idx + 1;
-    blockid_t jump_block = { ctx->iseq, jump_idx };
+    uint32_t next_idx = ctx_next_idx(ctx);
+    uint32_t jump_idx = next_idx + (uint32_t)ctx_get_arg(ctx, 0);
     blockid_t next_block = { ctx->iseq, next_idx };
+    blockid_t jump_block = { ctx->iseq, jump_idx };
 
     // TODO: we need to eventually do an interrupt check when jumping/branching
     // How can we do this while keeping the check logic out of line?
+    // Maybe we can push the check int into the next block or the stub?
+    //
 	// RUBY_VM_CHECK_INTS(ec);
 
     // Test if any bit (outside of the Qnil bit) is on
@@ -949,4 +947,5 @@ ujit_init_codegen(void)
     st_insert(gen_fns, (st_data_t)BIN(opt_minus), (st_data_t)&gen_opt_minus);
     st_insert(gen_fns, (st_data_t)BIN(opt_plus), (st_data_t)&gen_opt_plus);
     st_insert(gen_fns, (st_data_t)BIN(opt_send_without_block), (st_data_t)&gen_opt_send_without_block);
+    st_insert(gen_fns, (st_data_t)BIN(branchunless), (st_data_t)&gen_branchunless);
 }
diff --git a/ujit_core.c b/ujit_core.c
index 4a9ddc334d..c766f426a4 100644
--- a/ujit_core.c
+++ b/ujit_core.c
@@ -1,5 +1,10 @@
-#include "internal.h"
+#include "vm_core.h"
+#include "vm_callinfo.h"
+#include "builtin.h"
+#include "insns.inc"
+#include "insns_info.inc"
 #include "ujit_asm.h"
+#include "ujit_utils.h"
 #include "ujit_iface.h"
 #include "ujit_core.h"
 #include "ujit_codegen.h"
@@ -21,6 +26,13 @@ ctx_get_opcode(ctx_t *ctx)
     return opcode_at_pc(ctx->iseq, ctx->pc);
 }
 
+// Get the index of the next instruction
+uint32_t
+ctx_next_idx(ctx_t* ctx)
+{
+    return ctx->insn_idx + insn_len(ctx_get_opcode(ctx));
+}
+
 // Get an instruction argument from the context object
 VALUE
 ctx_get_arg(ctx_t* ctx, size_t arg_idx)
@@ -119,40 +131,46 @@ uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
 {
     assert (branch_idx < num_branches);
     assert (target_idx < 2);
-    branch_t branch = branch_entries[branch_idx];
-    blockid_t target = branch.targets[target_idx];
+    branch_t *branch = &branch_entries[branch_idx];
+    blockid_t target = branch->targets[target_idx];
+
+    //fprintf(stderr, "\nstub hit, branch idx: %d, target idx: %d\n", branch_idx, target_idx);
 
     // If either of the target blocks will be placed next
-    if (cb->write_pos == branch.end_pos)
+    if (cb->write_pos == branch->end_pos)
     {
-        branch.shape = (uint8_t)target_idx;
+        branch->shape = (uint8_t)target_idx;
 
         // Rewrite the branch with the new, potentially more compact shape
-        cb_set_pos(cb, branch.start_pos);
-        branch.gen_fn(cb, branch.dst_addrs[0], branch.dst_addrs[1], branch.shape);
-        assert (cb->write_pos <= branch.end_pos);
+        cb_set_pos(cb, branch->start_pos);
+        branch->gen_fn(cb, branch->dst_addrs[0], branch->dst_addrs[1], branch->shape);
+        assert (cb->write_pos <= branch->end_pos);
     }
 
     // Try to find a compiled version of this block
-    uint8_t* code_ptr = find_block_version(target);
+    uint8_t* block_ptr = find_block_version(target);
 
     // If this block hasn't yet been compiled
-    if (!code_ptr)
+    if (!block_ptr)
     {
-        code_ptr = ujit_compile_block(target.iseq, target.idx, false);
-        st_insert(version_tbl, (st_data_t)&target, (st_data_t)code_ptr);
-        branch.dst_addrs[target_idx] = code_ptr;
+        //fprintf(stderr, "compiling block\n");
+
+        block_ptr = ujit_compile_block(target.iseq, target.idx, false);
+        st_insert(version_tbl, (st_data_t)&target, (st_data_t)block_ptr);
+        branch->dst_addrs[target_idx] = block_ptr;
     }
 
+    //fprintf(stderr, "rewrite branch at %d\n", branch->start_pos);
+
     // Rewrite the branch with the new jump target address
     size_t cur_pos = cb->write_pos;
-    cb_set_pos(cb, branch.start_pos);
-    branch.gen_fn(cb, branch.dst_addrs[0], branch.dst_addrs[1], branch.shape);
-    assert (cb->write_pos <= branch.end_pos);
+    cb_set_pos(cb, branch->start_pos);
+    branch->gen_fn(cb, branch->dst_addrs[0], branch->dst_addrs[1], branch->shape);
+    assert (cb->write_pos <= branch->end_pos);
     cb_set_pos(cb, cur_pos);
 
     // Return a pointer to the compiled block version
-    return code_ptr;
+    return block_ptr;
 }
 
 // Get a version or stub corresponding to a branch target
@@ -164,14 +182,28 @@ uint8_t* get_branch_target(codeblock_t* ocb, blockid_t target, uint32_t branch_i
     if (block_code)
         return block_code;
 
-    uint8_t* stub_addr = cb_get_ptr(ocb, ocb->write_pos);
-
     // Generate an outlined stub that will call
     // branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
+    uint8_t* stub_addr = cb_get_ptr(ocb, ocb->write_pos);
+
+    //fprintf(stderr, "REQUESTING STUB FOR IDX: %d\n", target.idx);
+
+    // Save the ujit registers
+    push(ocb, REG_CFP);
+    push(ocb, REG_EC);
+    push(ocb, REG_SP);
+    push(ocb, REG_SP);
+
     mov(ocb, RDI, imm_opnd(branch_idx));
     mov(ocb, RSI, imm_opnd(target_idx));
     call_ptr(ocb, REG0, (void *)&branch_stub_hit);
 
+    // Restore the ujit registers
+    pop(ocb, REG_SP);
+    pop(ocb, REG_SP);
+    pop(ocb, REG_EC);
+    pop(ocb, REG_CFP);
+
     // Jump to the address returned by the
     // branch_stub_hit call
     jmp_rm(ocb, RAX);
diff --git a/ujit_core.h b/ujit_core.h
index c75b90be0e..6b9eaaaeac 100644
--- a/ujit_core.h
+++ b/ujit_core.h
@@ -98,6 +98,7 @@ typedef struct BranchEntry
 
 // Context object methods
 int ctx_get_opcode(ctx_t *ctx);
+uint32_t ctx_next_idx(ctx_t* ctx);
 VALUE ctx_get_arg(ctx_t* ctx, size_t arg_idx);
 x86opnd_t ctx_sp_opnd(ctx_t* ctx, int32_t offset_bytes);
 x86opnd_t ctx_stack_push(ctx_t* ctx, size_t n);