YJIT: adjust branch shape properly when target already exists

The old code decides branch->shape based on the write position of the
native code block, which is unsound in case the block already exists
and no new code is written to the write position.

Make this decision with the start address of the target block instead.
Also handle when the branch becomes smaller after patching.

2 files changed, 17 insertions, 16 deletions

diff --git a/yjit_core.c b/yjit_core.c
index ec793a053a..70c7da3291 100644
--- a/yjit_core.c
+++ b/yjit_core.c
@@ -418,7 +418,7 @@ uint8_t* gen_entry_point(const rb_iseq_t *iseq, uint32_t insn_idx, rb_execution_
 // Called by the generated code when a branch stub is executed
 // Triggers compilation of branches and code patching
 static uint8_t *
-branch_stub_hit(uint32_t branch_idx, uint32_t target_idx, rb_execution_context_t* ec)
+branch_stub_hit(const uint32_t branch_idx, const uint32_t target_idx, rb_execution_context_t* ec)
 {
     uint8_t* dst_addr;
 
@@ -449,18 +449,6 @@ branch_stub_hit(uint32_t branch_idx, uint32_t target_idx, rb_execution_context_t
     // 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)
-    {
-        //fprintf(stderr, "target idx %d will be placed next\n", target_idx);
-        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);
-        RUBY_ASSERT(cb->write_pos <= branch->end_pos);
-    }
-
     // Try to find an existing compiled version of this block
     block_t* p_block = find_block_version(target, target_ctx);
 
@@ -487,13 +475,26 @@ branch_stub_hit(uint32_t branch_idx, uint32_t target_idx, rb_execution_context_t
     dst_addr = cb_get_ptr(cb, p_block->start_pos);
     branch->dst_addrs[target_idx] = dst_addr;
 
+    // Adjust brach shape base on block placement relative to the branch
+    if (branch->end_pos == p_block->start_pos) {
+        branch->shape = (branch_shape_t)target_idx;
+    }
+
     // Rewrite the branch with the new jump target address
     RUBY_ASSERT(branch->dst_addrs[0] != NULL);
     uint32_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);
-    RUBY_ASSERT(cb->write_pos <= branch->end_pos);
-    branch->end_pos = cb->write_pos;
+    RUBY_ASSERT(cb->write_pos <= branch->end_pos && "can't enlarge a branch");
+
+    // If the branch got smaller
+    if (cb->write_pos < branch->end_pos) {
+        // fill the difference with nops
+        uint32_t shrinkage = branch->end_pos - cb->write_pos;
+        nop(cb, shrinkage);
+    }
+
+    // Done patching the branch. Restore write position.
     cb_set_pos(cb, cur_pos);
 
     // Restore interpreter sp, since the code hitting the stub expects the original.
diff --git a/yjit_core.h b/yjit_core.h
index 1a55408a7b..658a0cc691 100644
--- a/yjit_core.h
+++ b/yjit_core.h
@@ -156,7 +156,7 @@ typedef void (*branchgen_fn)(codeblock_t* cb, uint8_t* target0, uint8_t* target1
 Store info about an outgoing branch in a code segment
 Note: care must be taken to minimize the size of branch_t objects
 */
-typedef struct BranchEntry
+typedef struct yjit_branch_entry
 {
     // Positions where the generated code starts and ends
     uint32_t start_pos;