Refactor ujit logic for generating iseq entry points

ujit_codegen.c

@@ -48,18 +48,6 @@ jit_get_arg(jitstate_t* jit, size_t arg_idx)
     return *(jit->pc + arg_idx + 1);
 }
 
-/**
-Generate code to enter from the Ruby interpreter into ujit code
-*/
-static void
-ujit_gen_entry(codeblock_t* cb)
-{
-    cb_write_pre_call_bytes(cb);
-
-    // Load the current SP from the CFP into REG_SP
-    mov(cb, REG_SP, member_opnd(REG_CFP, rb_control_frame_t, sp));
-}
-
 /**
 Generate an inline exit to return to the interpreter
 */
@@ -110,10 +98,11 @@ ujit_side_exit(jitstate_t* jit, ctx_t* ctx)
 }
 
 /*
-Compile an interpreter entry point to be inserted into an iseq
+Compile an interpreter entry block to be inserted into an iseq
 Returns `NULL` if compilation fails.
 */
-uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx)
+uint8_t*
+ujit_gen_entry(version_t* version)
 {
     assert (cb != NULL);
 
@@ -127,36 +116,37 @@ uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx)
     uint8_t *code_ptr = cb_get_ptr(cb, cb->write_pos);
 
     // Write the interpreter entry prologue
-    ujit_gen_entry(cb);
+    cb_write_pre_call_bytes(cb);
 
-    // Create codegen context
-    ctx_t ctx = { 0 };
+    // Load the current SP from the CFP into REG_SP
+    mov(cb, REG_SP, member_opnd(REG_CFP, rb_control_frame_t, sp));
 
     // Compile the block starting at this instruction
-    uint32_t num_instrs = ujit_compile_block(iseq, insn_idx, &ctx);
+    uint32_t num_instrs = ujit_gen_code(version);
 
+    // FIXME: can we eliminate this check?
     // If no instructions were compiled
     if (num_instrs == 0) {
         return NULL;
     }
 
-    // Get the first opcode in the sequence
-    VALUE *encoded = iseq->body->iseq_encoded;
-    int first_opcode = opcode_at_pc(iseq, &encoded[insn_idx]);
-
-    // Map the code address to the corresponding opcode
-    map_addr2insn(code_ptr, first_opcode);
-
     return code_ptr;
 }
 
 /*
-Compile a sequence of bytecode instructions starting at `insn_idx`.
+Compile a sequence of bytecode instructions
 */
 uint32_t
-ujit_compile_block(/*version_t* version,*/ const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx)
+ujit_gen_code(version_t* version)
 {
     assert (cb != NULL);
+
+    // Copy the version's context to avoid mutating it
+    ctx_t ctx_copy = version->ctx;
+    ctx_t* ctx = &ctx_copy;
+
+    const rb_iseq_t *iseq = version->blockid.iseq;
+    uint32_t insn_idx = version->blockid.idx;
     VALUE *encoded = iseq->body->iseq_encoded;
 
     // NOTE: if we are ever deployed in production, we
@@ -174,9 +164,8 @@ ujit_compile_block(/*version_t* version,*/ const rb_iseq_t *iseq, uint32_t insn_
 
     // Initialize JIT state object
     jitstate_t jit = {
-        NULL,
-        iseq,
-        insn_idx
+        version,
+        iseq
     };
 
     uint32_t num_instrs = 0;
@@ -219,7 +208,7 @@ ujit_compile_block(/*version_t* version,*/ const rb_iseq_t *iseq, uint32_t insn_
         }
     }
 
-    // If the last instruction compiled did not properly terminate the block
+    // If the last instruction compiled did not terminate the block
     // Generate code to exit to the interpreter
     if (!p_last_op || !p_last_op->is_branch) {
         ujit_gen_exit(&jit, ctx, cb, &encoded[insn_idx]);
@@ -228,7 +217,7 @@ ujit_compile_block(/*version_t* version,*/ const rb_iseq_t *iseq, uint32_t insn_
     if (UJIT_DUMP_MODE >= 2) {
         // Dump list of compiled instrutions
         fprintf(stderr, "Compiled the following for iseq=%p:\n", (void *)iseq);
-        VALUE *pc = &encoded[jit.start_idx];
+        VALUE *pc = &encoded[version->blockid.idx];
         VALUE *end_pc = &encoded[insn_idx];
         while (pc < end_pc) {
             int opcode = opcode_at_pc(iseq, pc);

ujit_codegen.h

@@ -17,9 +17,6 @@ typedef struct JITState
     // Instruction sequence this is associated with
     const rb_iseq_t *iseq;
 
-    // Index in the iseq of the opcode we are replacing
-    const uint32_t start_idx;
-
     // Index of the current instruction being compiled
     uint32_t insn_idx;
 
@@ -43,9 +40,9 @@ typedef struct OpDesc
 
 } opdesc_t;
 
-uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx);
+uint8_t* ujit_gen_entry(version_t* version);
 
-uint32_t ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx);
+uint32_t ujit_gen_code(version_t* version);
 
 void ujit_init_codegen(void);
 

ujit_core.c

@@ -84,7 +84,6 @@ version_t* find_block_version(blockid_t block, const ctx_t* ctx)
 
     return NULL;
 }
-
 // Compile a new block version immediately
 version_t* gen_block_version(blockid_t blockid, const ctx_t* ctx)
 {
@@ -94,9 +93,8 @@ version_t* gen_block_version(blockid_t blockid, const ctx_t* ctx)
     memcpy(&p_version->ctx, ctx, sizeof(ctx_t));
 
     // Compile the block version
-    ctx_t ctx_copy = *ctx;
     p_version->start_pos = cb->write_pos;
-    ujit_compile_block(blockid.iseq, blockid.idx, &ctx_copy);
+    ujit_gen_code(p_version);
     p_version->end_pos = cb->write_pos;
 
     // Keep track of the new block version
@@ -105,6 +103,29 @@ version_t* gen_block_version(blockid_t blockid, const ctx_t* ctx)
     return p_version;
 }
 
+// Generate a block version that is an entry point inserted into an iseq
+uint8_t* gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx)
+{
+    // Allocate a version object
+    version_t* p_version = malloc(sizeof(version_t));
+    blockid_t blockid = { iseq, insn_idx };
+    memcpy(&p_version->blockid, &blockid, sizeof(blockid_t));
+
+    // The entry context makes no assumptions about types
+    ctx_t ctx = { 0 };
+    memcpy(&p_version->ctx, &ctx, sizeof(ctx_t));
+
+    // Compile the block version
+    p_version->start_pos = cb->write_pos;
+    uint8_t* code_ptr = ujit_gen_entry(p_version);
+    p_version->end_pos = cb->write_pos;
+
+    // Keep track of the new block version
+    st_insert(version_tbl, (st_data_t)&p_version->blockid, (st_data_t)p_version);
+
+    return code_ptr;
+}
+
 // Called by the generated code when a branch stub is executed
 // Triggers compilation of branches and code patching
 uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)

ujit_core.h

@@ -112,6 +112,7 @@ x86opnd_t ctx_stack_opnd(ctx_t* ctx, int32_t idx);
 
 version_t* find_block_version(blockid_t block, const ctx_t* ctx);
 version_t* gen_block_version(blockid_t block, const ctx_t* ctx);
+uint8_t*  gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx);
 
 void gen_branch(
     const ctx_t* src_ctx,

ujit_iface.c

@@ -237,10 +237,14 @@ rb_ujit_compile_iseq(const rb_iseq_t *iseq)
     VALUE *encoded = (VALUE *)iseq->body->iseq_encoded;
 
     // Compile a block version starting at the first instruction
-    uint8_t* native_code_ptr = ujit_compile_entry(iseq, 0);
+    uint8_t* code_ptr = gen_entry_point(iseq, 0);
 
-    if (native_code_ptr) {
-        encoded[0] = (VALUE)native_code_ptr;
+    if (code_ptr)
+    {
+        // Map the code address to the corresponding opcode
+        int first_opcode = opcode_at_pc(iseq, &encoded[0]);
+        map_addr2insn(code_ptr, first_opcode);
+        encoded[0] = (VALUE)code_ptr;
     }
 
     RB_VM_LOCK_LEAVE();