Complete refactoring to eliminate recursion in ujit's compilation

ujit_codegen.c

@@ -81,6 +81,9 @@ ujit_side_exit(jitstate_t* jit, ctx_t* ctx)
     // Table mapping opcodes to interpreter handlers
     const void * const *handler_table = rb_vm_get_insns_address_table();
 
+    // FIXME: rewriting the old instruction is only necessary if we're
+    // exiting right at an interpreter entry point
+
     // 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
@@ -102,7 +105,7 @@ Compile an interpreter entry block to be inserted into an iseq
 Returns `NULL` if compilation fails.
 */
 uint8_t*
-ujit_gen_entry(block_t* block)
+ujit_entry_prologue()
 {
     assert (cb != NULL);
 
@@ -121,28 +124,17 @@ ujit_gen_entry(block_t* block)
     // 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_gen_code(block);
-
-    // If no instructions were compiled
-    if (num_instrs == 0) {
-        return NULL;
-    }
-
     return code_ptr;
 }
 
 /*
 Compile a sequence of bytecode instructions for a given basic block version
 */
-uint32_t
+void
-ujit_gen_code(block_t* block)
+ujit_gen_block(ctx_t* ctx, block_t* block)
 {
     assert (cb != NULL);
-
+    assert (block != NULL);
-    // Copy the block's context to avoid mutating it
-    ctx_t ctx_copy = block->ctx;
-    ctx_t* ctx = &ctx_copy;
 
     const rb_iseq_t *iseq = block->blockid.iseq;
     uint32_t insn_idx = block->blockid.idx;
@@ -158,16 +150,19 @@ ujit_gen_code(block_t* block)
         rb_bug("out of executable memory (outlined block)");
     }
 
+    // Initialize a JIT state object
+    jitstate_t jit = {
+        block,
+        block->blockid.iseq,
+        0,
+        0
+    };
+
     // Last operation that was successfully compiled
     opdesc_t* p_last_op = NULL;
 
-    // Initialize JIT state object
+    // Mark the start position of the block
-    jitstate_t jit = {
+    block->start_pos = cb->write_pos;
-        block,
-        iseq
-    };
-
-    uint32_t num_instrs = 0;
 
     // For each instruction to compile
     for (;;) {
@@ -191,7 +186,7 @@ ujit_gen_code(block_t* block)
         opdesc_t* p_desc = (opdesc_t*)st_op_desc;
         bool success = p_desc->gen_fn(&jit, ctx);
 
-        // If we can't compile this instruction
+        // If we can't compile this instruction, stop
         if (!success) {
             break;
         }
@@ -199,7 +194,6 @@ ujit_gen_code(block_t* block)
     	// Move to the next instruction
         p_last_op = p_desc;
         insn_idx += insn_len(opcode);
-        num_instrs++;
 
         // If this instruction terminates this block
         if (p_desc->is_branch) {
@@ -207,15 +201,18 @@ ujit_gen_code(block_t* block)
         }
     }
 
-    // Store the index of the last instruction in the block
-    block->end_idx = insn_idx;
-
     // 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]);
     }
 
+    // Mark the end position of the block
+    block->end_pos = cb->write_pos;
+
+    // Store the index of the last instruction in the block
+    block->end_idx = insn_idx;
+
     if (UJIT_DUMP_MODE >= 2) {
         // Dump list of compiled instrutions
         fprintf(stderr, "Compiled the following for iseq=%p:\n", (void *)iseq);
@@ -227,8 +224,6 @@ ujit_gen_code(block_t* block)
             pc += insn_len(opcode);
         }
     }
-
-    return num_instrs;
 }
 
 static bool
@@ -683,7 +678,6 @@ gen_branchunless(jitstate_t* jit, ctx_t* ctx)
 
     // Generate the branch instructions
     gen_branch(
-        jit->block,
         ctx,
         jump_block,
         ctx,
@@ -727,7 +721,6 @@ gen_jump(jitstate_t* jit, ctx_t* ctx)
 
     // Generate the jump instruction
     gen_branch(
-        jit->block,
         ctx,
         jump_block,
         ctx,
@@ -986,7 +979,6 @@ gen_opt_send_without_block(jitstate_t* jit, ctx_t* ctx)
     // We do this to end the current block after the call
     blockid_t cont_block = { jit->iseq, jit_next_idx(jit) };
     gen_branch(
-        jit->block,
         ctx,
         cont_block,
         ctx,

ujit_codegen.h

@@ -40,9 +40,9 @@ typedef struct OpDesc
 
 } opdesc_t;
 
-uint8_t* ujit_gen_entry(block_t* block);
+uint8_t* ujit_entry_prologue();
 
-uint32_t ujit_gen_code(block_t* block);
+void ujit_gen_block(ctx_t* ctx, block_t* block);
 
 void ujit_init_codegen(void);
 

ujit_core.c

@@ -69,6 +69,17 @@ ctx_stack_opnd(ctx_t* ctx, int32_t idx)
     return opnd;
 }
 
+// Add an incoming branch for a given block version
+static void add_incoming(block_t* p_block, uint32_t branch_idx)
+{
+    // Add this branch to the list of incoming branches for the target
+    uint32_t* new_list = malloc(sizeof(uint32_t) * p_block->num_incoming + 1);
+    memcpy(new_list, p_block->incoming, p_block->num_incoming);
+    new_list[p_block->num_incoming] = branch_idx;
+    p_block->incoming = new_list;
+    p_block->num_incoming += 1;
+}
+
 // Retrieve a basic block version for an (iseq, idx) tuple
 block_t* find_block_version(blockid_t blockid, const ctx_t* ctx)
 {
@@ -79,87 +90,94 @@ block_t* find_block_version(blockid_t blockid, const ctx_t* ctx)
     }
 
     //
-    // TODO: use the ctx parameter to search available versions
+    // TODO: use the ctx parameter to search existing versions for a match
     //
 
     return NULL;
 }
 // Compile a new block version immediately
-block_t* gen_block_version(blockid_t blockid, const ctx_t* ctx)
+block_t* gen_block_version(blockid_t blockid, const ctx_t* start_ctx)
 {
-    // Allocate a version object
+    // Copy the context to avoid mutating it
-    block_t* p_block = malloc(sizeof(block_t));
+    ctx_t ctx_copy = *start_ctx;
-    memcpy(&p_block->blockid, &blockid, sizeof(blockid_t));
+    ctx_t* ctx = &ctx_copy;
-    memcpy(&p_block->ctx, ctx, sizeof(ctx_t));
-    p_block->incoming = NULL;
-    p_block->num_incoming = 0;
-    p_block->end_pos = 0;
 
-    // The block starts at the current position
+    // Allocate a new block version object
-    p_block->start_pos = cb->write_pos;
+    block_t* first_block = calloc(1, sizeof(block_t));
+    first_block->blockid = blockid;
+    memcpy(&first_block->ctx, ctx, sizeof(ctx_t));
 
-    // Compile the block version
+    // Block that is currently being compiled
-    ujit_gen_code(p_block);
+    block_t* block = first_block;
 
-    // The block may have been terminated in gen_branch
+    // Generate code for the first block
-    if (p_block->end_pos == 0)
+    ujit_gen_block(ctx, block);
-        p_block->end_pos = cb->write_pos;
 
     // Keep track of the new block version
-    st_insert(version_tbl, (st_data_t)&p_block->blockid, (st_data_t)p_block);
+    st_insert(version_tbl, (st_data_t)&block->blockid, (st_data_t)block);
 
-    return p_block;
+    // For each successor block to compile
+    for (;;) {
+        // If no branches were generated, stop
+        if (num_branches == 0) {
+            break;
+        }
+
+        // Get the last branch entry
+        uint32_t branch_idx = num_branches - 1;
+        branch_t* last_branch = &branch_entries[num_branches - 1];
+
+        // If there is no next block to compile, stop
+        if (last_branch->dst_addrs[0] || last_branch->dst_addrs[1]) {
+            break;
+        }
+
+        if (last_branch->targets[0].iseq == NULL) {
+            rb_bug("invalid target for last branch");
+        }
+
+        // Allocate a new block version object
+        block = calloc(1, sizeof(block_t));
+        block->blockid = last_branch->targets[0];
+        memcpy(&block->ctx, ctx, sizeof(ctx_t));
+
+        // Generate code for the current block
+        ujit_gen_block(ctx, block);
+
+        // Keep track of the new block version
+        st_insert(version_tbl, (st_data_t)&block->blockid, (st_data_t)block);
+
+        // Patch the last branch address
+        last_branch->dst_addrs[0] = cb_get_ptr(cb, block->start_pos);
+        add_incoming(block, branch_idx);
+        assert (block->start_pos == last_branch->end_pos);
+    }
+
+    return first_block;
 }
 
 // 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
-    block_t* p_block = malloc(sizeof(block_t));
-    blockid_t blockid = { iseq, insn_idx };
-    memcpy(&p_block->blockid, &blockid, sizeof(blockid_t));
-    p_block->incoming = NULL;
-    p_block->num_incoming = 0;
-    p_block->end_pos = 0;
-
     // The entry context makes no assumptions about types
+    blockid_t blockid = { iseq, insn_idx };
     ctx_t ctx = { 0 };
-    memcpy(&p_block->ctx, &ctx, sizeof(ctx_t));
 
-    // The block starts at the current position
+    // Write the interpreter entry prologue
-    p_block->start_pos = cb->write_pos;
+    uint8_t* code_ptr = ujit_entry_prologue();
 
-    // Compile the block version
+    // Try to generate code for the entry block
-    uint8_t* code_ptr = ujit_gen_entry(p_block);
+    block_t* block = gen_block_version(blockid, &ctx);
-
-    // The block may have been terminated in gen_branch
-    if (p_block->end_pos == 0)
-        p_block->end_pos = cb->write_pos;
 
     // If we couldn't generate any code
-    if (!code_ptr)
+    if (block->end_idx == insn_idx)
     {
-        free(p_block);
         return NULL;
     }
 
-    // Keep track of the new block version
-    st_insert(version_tbl, (st_data_t)&p_block->blockid, (st_data_t)p_block);
-
     return code_ptr;
 }
 
-// Add an incoming branch for a given block version
-static void add_incoming(block_t* p_block, uint32_t branch_idx)
-{
-    // Add this branch to the list of incoming branches for the target
-    uint32_t* new_list = malloc(sizeof(uint32_t) * p_block->num_incoming + 1);
-    memcpy(new_list, p_block->incoming, p_block->num_incoming);
-    new_list[p_block->num_incoming] = branch_idx;
-    p_block->incoming = new_list;
-    p_block->num_incoming += 1;
-}
-
 // 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)
@@ -265,7 +283,6 @@ uint8_t* get_branch_target(
 }
 
 void gen_branch(
-    block_t* src_version,
     const ctx_t* src_ctx,
     blockid_t target0, 
     const ctx_t* ctx0,
@@ -274,42 +291,50 @@ void gen_branch(
     branchgen_fn gen_fn
 )
 {
+    assert (target0.iseq != NULL);
     assert (num_branches < MAX_BRANCHES);
-    uint32_t branch_idx = num_branches;
+    uint32_t branch_idx = num_branches++;
 
     // Branch targets or stub adddresses (code pointers)
     uint8_t* dst_addr0;
     uint8_t* dst_addr1;
 
+    // Shape of the branch
+    uint8_t branch_shape;
+
     // If there's only one branch target
     if (target1.iseq == NULL)
     {
         block_t* p_block = find_block_version(target0, ctx0);
 
-        // If the version doesn't already exist
+        // If the version already exists
-        if (!p_block)
+        if (p_block)
         {
-            // No need for a jump, compile the target block right here
-            p_block = gen_block_version(target0, ctx0);
-
-            // The current version ends where the next version begins
-            src_version->end_pos = p_block->start_pos;
-        }
-
             add_incoming(p_block, branch_idx);
             dst_addr0 = cb_get_ptr(cb, p_block->start_pos);
             dst_addr1 = NULL;
+            branch_shape = SHAPE_DEFAULT;
+        }
+        else
+        {
+            // The target block will follow next
+            // It will be compiled in gen_block_version()
+            dst_addr0 = NULL;
+            dst_addr1 = NULL;
+            branch_shape = SHAPE_NEXT0;
+        }
     }
     else
     {
         // Get the branch targets or stubs
         dst_addr0 = get_branch_target(target0, ctx0, branch_idx, 0);
         dst_addr1 = get_branch_target(target1, ctx1, branch_idx, 1);
+        branch_shape = SHAPE_DEFAULT;
     }
 
     // Call the branch generation function
     uint32_t start_pos = cb->write_pos;
-    gen_fn(cb, dst_addr0, dst_addr1, SHAPE_DEFAULT);
+    gen_fn(cb, dst_addr0, dst_addr1, branch_shape);
     uint32_t end_pos = cb->write_pos;
 
     // Register this branch entry
@@ -321,17 +346,16 @@ void gen_branch(
         { *ctx0, *ctx1 },
         { dst_addr0, dst_addr1 },
         gen_fn,
-        SHAPE_DEFAULT
+        branch_shape
     };
 
-    branch_entries[num_branches] = branch_entry;
+    branch_entries[branch_idx] = branch_entry;
-    num_branches++;
 }
 
 // Invalidate one specific block version
 void invalidate(block_t* block)
 {
-    fprintf(stderr, "invalidata block (%p, %d)\n", block->blockid.iseq, block->blockid.idx);
+    fprintf(stderr, "invalidating block (%p, %d)\n", block->blockid.iseq, block->blockid.idx);
 
     // Remove the version object from the map so we can re-generate stubs
     st_delete(version_tbl, (st_data_t*)&block->blockid, NULL);

ujit_core.h

@@ -37,8 +37,8 @@ typedef struct BlockId
     // Instruction sequence
     const rb_iseq_t *iseq;
 
-    // Instruction index
+    // Index in the iseq where the block starts
-    const uint32_t idx;
+    uint32_t idx;
 
 } blockid_t;
 
@@ -87,7 +87,7 @@ typedef struct BlockVersion
     // Bytecode sequence (iseq, idx) this is a version of
     blockid_t blockid;
 
-    // Index just past the last instruction in the iseq
+    // Index one past the last instruction in the iseq
     uint32_t end_idx;
 
     // Context at the start of the block
@@ -117,7 +117,6 @@ block_t* gen_block_version(blockid_t blockid, const ctx_t* ctx);
 uint8_t*  gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx);
 
 void gen_branch(
-    block_t* src_block,
     const ctx_t* src_ctx,
     blockid_t target0,
     const ctx_t* ctx0,