Implement branch stub logic

ujit_codegen.c

@@ -20,11 +20,11 @@ static st_table *gen_fns;
 
 // Code block into which we write machine code
 static codeblock_t block;
-static codeblock_t* cb = NULL;
+codeblock_t* cb = NULL;
 
 // Code block into which we write out-of-line machine code
 static codeblock_t outline_block;
-static codeblock_t* ocb = NULL;
+codeblock_t* ocb = NULL;
 
 // Ruby instruction entry
 static void

ujit_codegen.h

@@ -3,6 +3,10 @@
 
 #include "stddef.h"
 
+// Code blocks we generate code into
+codeblock_t* cb;
+codeblock_t* ocb;
+
 // Code generation function signature
 typedef bool (*codegen_fn)(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx);
 

ujit_core.c

@@ -101,25 +101,6 @@ static const struct st_hash_type hashtype_blockid = {
     blockid_hash,
 };
 
-// Called by the generated code when a branch stub is executed
-// Triggers compilation of branches and code patching
-void branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
-{
-
-
-
-
-    // TODO
-    //uint8_t* code_ptr = ujit_compile_block(blockid.iseq, blockid.idx, false);
-    //st_insert(version_tbl, (st_data_t)&blockid, (st_data_t)code_ptr);
-
-
-
-
-
-
-}
-
 // Retrieve a basic block version for an (iseq, idx) tuple
 uint8_t* find_block_version(blockid_t block)
 {
@@ -132,6 +113,48 @@ uint8_t* find_block_version(blockid_t block)
     return NULL;
 }
 
+// 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)
+{
+    assert (branch_idx < num_branches);
+    assert (target_idx < 2);
+    branch_t branch = branch_entries[branch_idx];
+    blockid_t target = branch.targets[target_idx];
+
+    // If either of the target blocks will be placed next
+    if (cb->write_pos == branch.end_pos)
+    {
+        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);
+    }
+
+    // Try to find a compiled version of this block
+    uint8_t* code_ptr = find_block_version(target);
+
+    // If this block hasn't yet been compiled
+    if (!code_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;
+    }
+
+    // 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, cur_pos);
+
+    // Return a pointer to the compiled block version
+    return code_ptr;
+}
+
 // Get a version or stub corresponding to a branch target
 // TODO: need incoming and target versioning contexts
 uint8_t* get_branch_target(codeblock_t* ocb, blockid_t target, uint32_t branch_idx, uint32_t target_idx)
@@ -145,16 +168,13 @@ uint8_t* get_branch_target(codeblock_t* ocb, blockid_t target, uint32_t branch_i
 
     // Generate an outlined stub that will call
     // branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
+    mov(ocb, RDI, imm_opnd(branch_idx));
+    mov(ocb, RSI, imm_opnd(target_idx));
+    call_ptr(ocb, REG0, (void *)&branch_stub_hit);
 
-
-
-
-
-
-
-
-
-
+    // Jump to the address returned by the
+    // branch_stub_hit call
+    jmp_rm(ocb, RAX);
 
     return stub_addr;
 }
@@ -162,13 +182,13 @@ uint8_t* get_branch_target(codeblock_t* ocb, blockid_t target, uint32_t branch_i
 void gen_branch(codeblock_t* cb, codeblock_t* ocb, blockid_t target0, blockid_t target1, branchgen_fn gen_fn)
 {
     // Get branch targets or stubs (code pointers)
-    uint8_t* target_code0 = get_branch_target(ocb, target0, num_branches, 0);
-    uint8_t* target_code1 = get_branch_target(ocb, target1, num_branches, 1);
+    uint8_t* dst_addr0 = get_branch_target(ocb, target0, num_branches, 0);
+    uint8_t* dst_addr1 = get_branch_target(ocb, target1, num_branches, 1);
 
     uint32_t start_pos = (uint32_t)cb->write_pos;
 
     // Call the branch generation function
-    gen_fn(cb, target_code0, target_code1, DEFAULT);
+    gen_fn(cb, dst_addr0, dst_addr1, SHAPE_DEFAULT);
 
     uint32_t end_pos = (uint32_t)cb->write_pos;
 
@@ -177,7 +197,9 @@ void gen_branch(codeblock_t* cb, codeblock_t* ocb, blockid_t target0, blockid_t
         start_pos,
         end_pos,
         { target0, target1 },
-        gen_fn
+        { dst_addr0, dst_addr1 },
+        gen_fn,
+        SHAPE_DEFAULT
     };
 
     assert (num_branches < MAX_BRANCHES);

ujit_core.h

@@ -67,9 +67,9 @@ typedef struct BlockId
 /// Branch code shape enumeration
 enum uint8_t
 {
-    NEXT0,  // Target 0 is next
-    NEXT1,  // Target 1 is next
-    DEFAULT // Neither target is next
+    SHAPE_NEXT0,  // Target 0 is next
+    SHAPE_NEXT1,  // Target 1 is next
+    SHAPE_DEFAULT // Neither target is next
 };
 
 // Branch code generation function signature
@@ -85,9 +85,15 @@ typedef struct BranchEntry
     // Branch target blocks
     blockid_t targets[2];
 
+    // Jump target addresses
+    uint8_t* dst_addrs[2];
+
     // Branch code generation function
     branchgen_fn gen_fn;
 
+    // Shape of the branch
+    uint8_t shape;
+
 } branch_t;
 
 // Context object methods
@@ -98,8 +104,6 @@ x86opnd_t ctx_stack_push(ctx_t* ctx, size_t n);
 x86opnd_t ctx_stack_pop(ctx_t* ctx, size_t n);
 x86opnd_t ctx_stack_opnd(ctx_t* ctx, int32_t idx);
 
-uint8_t* get_block_version(blockid_t block);
-
 void gen_branch(codeblock_t* cb, codeblock_t* ocb, blockid_t target0, blockid_t target1, branchgen_fn gen_fn);
 
 void ujit_init_core(void);