Update ec->cfp->pc when hitting a stub.

ujit_codegen.c

@@ -66,8 +66,10 @@ jit_mov_gc_ptr(jitstate_t* jit, codeblock_t* cb, x86opnd_t reg, VALUE ptr)
 {
     RUBY_ASSERT(reg.type == OPND_REG && reg.num_bits == 64);
 
+    // Load the pointer constant into the specified register
     mov(cb, reg, const_ptr_opnd((void*)ptr));
-    // The pointer immediate is encoded as the last part of the mov written out.
+
+    // The pointer immediate is encoded as the last part of the mov written out
     uint32_t ptr_offset = cb->write_pos - sizeof(VALUE);
 
     if (!SPECIAL_CONST_P(ptr)) {
@@ -147,7 +149,7 @@ ujit_side_exit(jitstate_t* jit, ctx_t* ctx)
     // 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
-    VALUE* exit_pc = &jit->iseq->body->iseq_encoded[jit->insn_idx];
+    VALUE* exit_pc = iseq_pc_at_idx(jit->iseq, jit->insn_idx);
     int exit_opcode = opcode_at_pc(jit->iseq, exit_pc);
     void* handler_addr = (void*)handler_table[exit_opcode];
     mov(ocb, RAX, const_ptr_opnd(exit_pc));
@@ -212,7 +214,6 @@ ujit_gen_block(ctx_t* ctx, block_t* block, rb_execution_context_t* ec)
 
     const rb_iseq_t *iseq = block->blockid.iseq;
     uint32_t insn_idx = block->blockid.idx;
-    VALUE *encoded = iseq->body->iseq_encoded;
 
     // NOTE: if we are ever deployed in production, we
     // should probably just log an error and return NULL here,
@@ -227,7 +228,7 @@ ujit_gen_block(ctx_t* ctx, block_t* block, rb_execution_context_t* ec)
     // Initialize a JIT state object
     jitstate_t jit = {
         block,
-        block->blockid.iseq,
+        iseq,
         0,
         0,
         ec
@@ -239,9 +240,8 @@ ujit_gen_block(ctx_t* ctx, block_t* block, rb_execution_context_t* ec)
     // For each instruction to compile
     for (;;) {
         // Set the current instruction
-        RUBY_ASSERT(insn_idx < iseq->body->iseq_size);
         jit.insn_idx = insn_idx;
-        jit.pc = &encoded[insn_idx];
+        jit.pc = iseq_pc_at_idx(iseq, insn_idx);
 
         // Get the current opcode
         int opcode = jit_get_opcode(&jit);
@@ -251,7 +251,7 @@ ujit_gen_block(ctx_t* ctx, block_t* block, rb_execution_context_t* ec)
         if (!rb_st_lookup(gen_fns, opcode, (st_data_t*)&gen_fn)) {
             // If we reach an unknown instruction,
             // exit to the interpreter and stop compiling
-            ujit_gen_exit(&jit, ctx, cb, &encoded[insn_idx]);
+            ujit_gen_exit(&jit, ctx, cb, jit.pc);
             break;
         }
 
@@ -273,7 +273,7 @@ ujit_gen_block(ctx_t* ctx, block_t* block, rb_execution_context_t* ec)
         // If we can't compile this instruction
         // exit to the interpreter and stop compiling
         if (status == UJIT_CANT_COMPILE) {
-            ujit_gen_exit(&jit, ctx, cb, &encoded[insn_idx]);
+            ujit_gen_exit(&jit, ctx, cb, jit.pc);
             break;
         }
 
@@ -295,12 +295,11 @@ ujit_gen_block(ctx_t* ctx, block_t* block, rb_execution_context_t* ec)
     if (UJIT_DUMP_MODE >= 2) {
         // Dump list of compiled instrutions
         fprintf(stderr, "Compiled the following for iseq=%p:\n", (void *)iseq);
-        VALUE *pc = &encoded[block->blockid.idx];
-        VALUE *end_pc = &encoded[insn_idx];
-        while (pc < end_pc) {
-            int opcode = opcode_at_pc(iseq, pc);
-            fprintf(stderr, "  %04td %s\n", pc - encoded, insn_name(opcode));
-            pc += insn_len(opcode);
+        for (uint32_t idx = block->blockid.idx; idx < insn_idx;)
+        {
+            int opcode = opcode_at_pc(iseq, iseq_pc_at_idx(iseq, idx));
+            fprintf(stderr, "  %04d %s\n", idx, insn_name(opcode));
+            idx += insn_len(opcode);
         }
     }
 }

ujit_core.c

@@ -373,6 +373,10 @@ uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx, rb_execution_
     //fprintf(stderr, "\nstub hit, branch idx: %d, target idx: %d\n", branch_idx, target_idx);
     //fprintf(stderr, "blockid.iseq=%p, blockid.idx=%d\n", target.iseq, target.idx);
 
+    // Update the PC in the current CFP, because it
+    // may be out of sync in JITted code
+    ec->cfp->pc = iseq_pc_at_idx(target.iseq, target.idx);
+
     // If either of the target blocks will be placed next
     if (cb->write_pos == branch->end_pos)
     {
@@ -688,7 +692,7 @@ invalidate_block_version(block_t* block)
     uint32_t idx = block->blockid.idx;
     // FIXME: the following says "if", but it's unconditional.
     // If the block is an entry point, it needs to be unmapped from its iseq
-    VALUE* entry_pc = &iseq->body->iseq_encoded[idx];
+    VALUE* entry_pc = iseq_pc_at_idx(iseq, idx);
     int entry_opcode = opcode_at_pc(iseq, entry_pc);
 
     // TODO: unmap_addr2insn in ujit_iface.c? Maybe we can write a function to encompass this logic?

ujit_iface.c

@@ -62,6 +62,16 @@ cb_write_post_call_bytes(codeblock_t* cb)
         cb_write_byte(cb, ujit_with_ec_post_call_bytes[i]);
 }
 
+// Get the PC for a given index in an iseq
+VALUE *iseq_pc_at_idx(const rb_iseq_t *iseq, uint32_t insn_idx)
+{
+    RUBY_ASSERT(iseq != NULL);
+    RUBY_ASSERT(insn_idx < iseq->body->iseq_size);
+    VALUE *encoded = iseq->body->iseq_encoded;
+    VALUE *pc = &encoded[insn_idx];
+    return pc;
+}
+
 // Keep track of mapping from instructions to generated code
 // See comment for rb_encoded_insn_data in iseq.c
 void

ujit_iface.h

@@ -28,6 +28,7 @@ RUBY_EXTERN int64_t rb_compiled_iseq_count;
 void cb_write_pre_call_bytes(codeblock_t* cb);
 void cb_write_post_call_bytes(codeblock_t* cb);
 
+VALUE *iseq_pc_at_idx(const rb_iseq_t *iseq, uint32_t insn_idx);
 void map_addr2insn(void *code_ptr, int insn);
 int opcode_at_pc(const rb_iseq_t *iseq, const VALUE *pc);