WIP branch generation code

4 files changed, 189 insertions, 60 deletions

diff --git a/ujit_codegen.c b/ujit_codegen.c
index 01785d5fc4..1201c99c66 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -15,9 +15,6 @@
 #include "ujit_asm.h"
 #include "ujit_utils.h"
 
-// Code generation function signature
-typedef bool (*codegen_fn)(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx);
-
 // Map from YARV opcodes to code generation functions
 static st_table *gen_fns;
 
@@ -889,6 +886,38 @@ gen_opt_send_without_block(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx)
     return true;
 }
 
+void 
+gen_branchunless_branch(codeblock_t* cb, uint8_t* target0, uint8_t* target1, uint8_t shape)
+{
+    jz_ptr(cb, target0);
+    jmp_ptr(cb, target1);
+}
+
+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 };
+    blockid_t next_block = { ctx->iseq, next_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?
+	// RUBY_VM_CHECK_INTS(ec);
+
+    // Test if any bit (outside of the Qnil bit) is on
+    // RUBY_Qfalse  /* ...0000 0000 */
+    // RUBY_Qnil    /* ...0000 1000 */
+    x86opnd_t val_opnd = ctx_stack_pop(ctx, 1);
+    test(cb, val_opnd, imm_opnd(~Qnil));
+
+    // Generate the branch instructions
+    gen_branch(cb, ocb, jump_block, next_block, gen_branchunless_branch);
+
+    return true;
+}
+
 void
 ujit_init_codegen(void)
 {
diff --git a/ujit_codegen.h b/ujit_codegen.h
index 3dc3d1bf47..7de90c7877 100644
--- a/ujit_codegen.h
+++ b/ujit_codegen.h
@@ -3,6 +3,9 @@
 
 #include "stddef.h"
 
+// Code generation function signature
+typedef bool (*codegen_fn)(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx);
+
 uint8_t *ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, bool gen_entry);
 
 void ujit_init_codegen(void);
diff --git a/ujit_core.c b/ujit_core.c
index 1119e3df91..916e8bc4d8 100644
--- a/ujit_core.c
+++ b/ujit_core.c
@@ -4,54 +4,15 @@
 #include "ujit_core.h"
 #include "ujit_codegen.h"
 
+// Maximum number of branch instructions we can track
+#define MAX_BRANCHES 32768
+
 // Table of block versions indexed by (iseq, index) tuples
 st_table * version_tbl;
 
-int blockid_cmp(st_data_t arg0, st_data_t arg1)
-{
-    const blockid_t *block0 = (const blockid_t*)arg0;
-    const blockid_t *block1 = (const blockid_t*)arg1;
-    return block0->iseq == block1->iseq && block0->idx == block1->idx;
-}
-
-st_index_t blockid_hash(st_data_t arg)
-{
-    const blockid_t *blockid = (const blockid_t*)arg;
-    st_index_t hash0 = st_numhash((st_data_t)blockid->iseq);
-    st_index_t hash1 = st_numhash((st_data_t)(uint64_t)blockid->idx);
-
-    // Use XOR to combine the hashes
-    return hash0 ^ hash1;
-}
-
-static const struct st_hash_type hashtype_blockid = {
-    blockid_cmp,
-    blockid_hash,
-};
-
-// Retrieve a basic block version for an (iseq, idx) tuple
-// TODO: we need to add a versioning context here
-uint8_t* get_block_version(const rb_iseq_t *iseq, uint32_t idx)
-{
-    blockid_t blockid = { iseq, idx };
-
-    // If there exists a version for this block id
-    st_data_t st_version;
-    if (rb_st_lookup(version_tbl, (st_data_t)&blockid, &st_version)) {
-        return (uint8_t*)st_version;
-    }
-
-    uint8_t* code_ptr = ujit_compile_block(iseq, idx, false);
-
-    st_insert(version_tbl, (st_data_t)&blockid, (st_data_t)code_ptr);
-
-    return code_ptr;
-}
-
-//
-// Method to generate stubs for branches
-// TODO: get_branch_stub() or get_branch() function
-//
+// Registered branch entries
+branch_t branch_entries[MAX_BRANCHES];
+uint32_t num_branches = 0;
 
 // Get the current instruction opcode from the context object
 int
@@ -118,6 +79,112 @@ ctx_stack_opnd(ctx_t* ctx, int32_t idx)
     return opnd;
 }
 
+int blockid_cmp(st_data_t arg0, st_data_t arg1)
+{
+    const blockid_t *block0 = (const blockid_t*)arg0;
+    const blockid_t *block1 = (const blockid_t*)arg1;
+    return block0->iseq == block1->iseq && block0->idx == block1->idx;
+}
+
+st_index_t blockid_hash(st_data_t arg)
+{
+    const blockid_t *blockid = (const blockid_t*)arg;
+    st_index_t hash0 = st_numhash((st_data_t)blockid->iseq);
+    st_index_t hash1 = st_numhash((st_data_t)(uint64_t)blockid->idx);
+
+    // Use XOR to combine the hashes
+    return hash0 ^ hash1;
+}
+
+static const struct st_hash_type hashtype_blockid = {
+    blockid_cmp,
+    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)
+{
+    // If there exists a version for this block id
+    st_data_t st_version;
+    if (rb_st_lookup(version_tbl, (st_data_t)&block, &st_version)) {
+        return (uint8_t*)st_version;
+    }
+
+    return NULL;
+}
+
+// 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)
+{
+    uint8_t* block_code = find_block_version(target);
+
+    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)
+
+
+
+
+
+
+
+
+
+
+
+    return stub_addr;
+}
+
+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);
+
+    uint32_t start_pos = (uint32_t)cb->write_pos;
+
+    // Call the branch generation function
+    gen_fn(cb, target_code0, target_code1, DEFAULT);
+
+    uint32_t end_pos = (uint32_t)cb->write_pos;
+
+    // Register this branch entry
+    branch_t branch_entry = {
+        start_pos,
+        end_pos,
+        { target0, target1 },
+        gen_fn
+    };
+
+    assert (num_branches < MAX_BRANCHES);
+    branch_entries[num_branches] = branch_entry;
+    num_branches++;
+}
+
 void
 ujit_init_core(void)
 {
diff --git a/ujit_core.h b/ujit_core.h
index 5bcbd89f3c..c6c43de525 100644
--- a/ujit_core.h
+++ b/ujit_core.h
@@ -20,17 +20,6 @@
 // Maximum number of versions per block
 #define MAX_VERSIONS 5
 
-// Tuple of (iseq, idx) used to idenfity basic blocks
-typedef struct BlockId
-{
-    // Instruction sequence
-    const rb_iseq_t *iseq;
-
-    // Instruction index
-    const uint32_t idx;
-
-} blockid_t;
-
 // Code generation context
 typedef struct ctx_struct
 {
@@ -39,6 +28,8 @@ typedef struct ctx_struct
     // Some of the information here is only needed during
     // code generation, eg: current pc
 
+    // FIXME: we probably don't need this? we just need to
+    // know which initial bytecode we're replacing
     // The start of the generated code
     uint8_t *code_ptr;
 
@@ -62,7 +53,42 @@ typedef struct ctx_struct
 
 } ctx_t;
 
-uint8_t* get_block_version(const rb_iseq_t *iseq, uint32_t idx);
+// Tuple of (iseq, idx) used to idenfity basic blocks
+typedef struct BlockId
+{
+    // Instruction sequence
+    const rb_iseq_t *iseq;
+
+    // Instruction index
+    const uint32_t idx;
+
+} blockid_t;
+
+/// Branch code shape enumeration
+enum uint8_t
+{
+    NEXT0,  // Target 0 is next
+    NEXT1,  // Target 1 is next
+    DEFAULT // Neither target is next
+};
+
+// Branch code generation function signature
+typedef void (*branchgen_fn)(codeblock_t* cb, uint8_t* target0, uint8_t* target1, uint8_t shape);
+
+// Store info about an outgoing branch in a code segment
+typedef struct BranchEntry
+{
+    // Positions where the generated code starts and ends
+    uint32_t start_pos;
+    uint32_t end_pos;
+
+    // Branch target blocks
+    blockid_t targets[2];
+
+    // Branch code generation function
+    branchgen_fn gen_fn;
+
+} branch_t;
 
 // Context object methods
 int ctx_get_opcode(ctx_t *ctx);
@@ -72,6 +98,10 @@ 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);
 
 #endif // #ifndef UJIT_CORE_H