3 files changed, 88 insertions, 65 deletions

diff --git a/ujit_codegen.c b/ujit_codegen.c
index bc91a9ad3b..8db6942de1 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -689,24 +689,6 @@ gen_branchunless(jitstate_t* jit, ctx_t* ctx)
     return true;
 }
 
-void
-gen_jump_branch(codeblock_t* cb, uint8_t* target0, uint8_t* target1, uint8_t shape)
-{
-    switch (shape)
-    {
-        case SHAPE_NEXT0:
-        break;
-
-        case SHAPE_NEXT1:
-        assert (false);
-        break;
-
-        case SHAPE_DEFAULT:
-        jmp_ptr(cb, target0);
-        break;
-    }
-}
-
 static bool
 gen_jump(jitstate_t* jit, ctx_t* ctx)
 {
@@ -720,13 +702,9 @@ gen_jump(jitstate_t* jit, ctx_t* ctx)
     //
 
     // Generate the jump instruction
-    gen_branch(
+    gen_direct_jump(
         ctx,
-        jump_block,
-        ctx,
-        BLOCKID_NULL,
-        ctx,
-        gen_jump_branch
+        jump_block
     );
 
     return true;
@@ -978,13 +956,9 @@ gen_opt_send_without_block(jitstate_t* jit, ctx_t* ctx)
     // Jump (fall through) to the call continuation block
     // We do this to end the current block after the call
     blockid_t cont_block = { jit->iseq, jit_next_idx(jit) };
-    gen_branch(
-        ctx,
-        cont_block,
-        ctx,
-        BLOCKID_NULL,
+    gen_direct_jump(
         ctx,
-        gen_jump_branch
+        cont_block
     );
 
     return true;
diff --git a/ujit_core.c b/ujit_core.c
index 332e356860..c6ba70bb28 100644
--- a/ujit_core.c
+++ b/ujit_core.c
@@ -292,60 +292,104 @@ void gen_branch(
 )
 {
     assert (target0.iseq != NULL);
+    assert (target1.iseq != NULL);
     assert (num_branches < MAX_BRANCHES);
     uint32_t branch_idx = num_branches++;
 
-    // Branch targets or stub adddresses (code pointers)
+    // Get the branch targets or stubs
+    uint8_t* dst_addr0 = get_branch_target(target0, ctx0, branch_idx, 0);
+    uint8_t* dst_addr1 = get_branch_target(target1, ctx1, branch_idx, 1);
+
+    // Call the branch generation function
+    uint32_t start_pos = cb->write_pos;
+    gen_fn(cb, dst_addr0, dst_addr1, SHAPE_DEFAULT);
+    uint32_t end_pos = cb->write_pos;
+
+    // Register this branch entry
+    branch_t branch_entry = {
+        start_pos,
+        end_pos,
+        *src_ctx,
+        { target0, target1 },
+        { *ctx0, *ctx1 },
+        { dst_addr0, dst_addr1 },
+        gen_fn,
+        SHAPE_DEFAULT
+    };
+
+    branch_entries[branch_idx] = branch_entry;
+}
+
+void
+gen_jump_branch(codeblock_t* cb, uint8_t* target0, uint8_t* target1, uint8_t shape)
+{
+    switch (shape)
+    {
+        case SHAPE_NEXT0:
+        break;
+
+        case SHAPE_NEXT1:
+        assert (false);
+        break;
+
+        case SHAPE_DEFAULT:
+        jmp_ptr(cb, target0);
+        break;
+    }
+}
+
+void gen_direct_jump(
+    const ctx_t* ctx,
+    blockid_t target0
+)
+{
+    assert (target0.iseq != NULL);
+    assert (num_branches < MAX_BRANCHES);
+    uint32_t branch_idx = num_branches++;
+
+    // Branch targets or stub adddress
     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)
+    // Branch start and end positions
+    uint32_t start_pos;
+    uint32_t end_pos;
+
+    block_t* p_block = find_block_version(target0, ctx);
+
+    // If the version already exists
+    if (p_block)
     {
-        block_t* p_block = find_block_version(target0, ctx0);
+        add_incoming(p_block, branch_idx);
+        dst_addr0 = cb_get_ptr(cb, p_block->start_pos);
+        branch_shape = SHAPE_DEFAULT;
 
-        // If the version already exists
-        if (p_block)
-        {
-            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;
-        }
+        // Call the branch generation function
+        start_pos = cb->write_pos;
+        gen_jump_branch(cb, dst_addr0, NULL, branch_shape);
+        end_pos = cb->write_pos;
     }
     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;
+        // The target block will follow next
+        // It will be compiled in gen_block_version()
+        dst_addr0 = NULL;
+        branch_shape = SHAPE_NEXT0;
+        start_pos = cb->write_pos;
+        end_pos = cb->write_pos;
     }
 
-    // Call the branch generation function
-    uint32_t start_pos = cb->write_pos;
-    gen_fn(cb, dst_addr0, dst_addr1, branch_shape);
-    uint32_t end_pos = cb->write_pos;
-
     // Register this branch entry
     branch_t branch_entry = {
         start_pos,
         end_pos,
-        *src_ctx,
-        { target0, target1 },
-        { *ctx0, *ctx1 },
-        { dst_addr0, dst_addr1 },
-        gen_fn,
+        *ctx,
+        { target0, BLOCKID_NULL },
+        { *ctx, *ctx },
+        { dst_addr0, NULL },
+        gen_jump_branch,
         branch_shape
     };
 
diff --git a/ujit_core.h b/ujit_core.h
index 3e119758e0..c99bd808a0 100644
--- a/ujit_core.h
+++ b/ujit_core.h
@@ -125,6 +125,11 @@ void gen_branch(
     branchgen_fn gen_fn
 );
 
+void gen_direct_jump(
+    const ctx_t* ctx,
+    blockid_t target0
+);
+
 void invalidate(block_t* block);
 
 void ujit_init_core(void);