From 4592ef9d76c5ef1bf7f99119a126a1b700c88f1f Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Thu, 17 Dec 2020 15:45:38 -0500
Subject: [PATCH] Fixed one bug in compile_block.

---
 ujit_codegen.c | 30 ++++++++++++++++++------------
 ujit_core.c    |  2 ++
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/ujit_codegen.c b/ujit_codegen.c
index 215bc3cf92..ebffef424a 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -145,6 +145,9 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, bool entry_point)
             ujit_gen_entry(cb);
         }
 
+        //fprintf(stderr, "compiling %s\n", insn_name(opcode));
+        //print_str(cb, insn_name(opcode));
+
         // Call the code generation function
         codegen_fn gen_fn = (codegen_fn)st_gen_fn;
         if (!gen_fn(cb, ocb, &ctx)) {
@@ -163,13 +166,16 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, bool entry_point)
         }
     }
 
+    // FIXME: maybe we want a separate function to compile entry points?
+    // If this is an entry point and no instructions were compiled
+    if (entry_point && num_instrs == 0) {
+        return NULL;
+    }
+
+    //print_str(cb, "exiting to interpreter\n");
+
     // FIXME: only generate exit if no instructions were compiled?
     // or simply don't allow instructions to fail to compile anymore?
-    // If no instructions were compiled
-    //if (num_instrs == 0) {
-    //    return NULL;
-    //}
-
     // Generate code to exit to the interpreter
     ujit_gen_exit(cb, &ctx, &encoded[insn_idx]);
 
@@ -892,12 +898,6 @@ gen_branchunless_branch(codeblock_t* cb, uint8_t* target0, uint8_t* target1, uin
 static bool
 gen_branchunless(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx)
 {
-    // Get the branch target instruction offsets
-    uint32_t next_idx = ctx_next_idx(ctx);
-    uint32_t jump_idx = next_idx + (uint32_t)ctx_get_arg(ctx, 0);
-    blockid_t next_block = { ctx->iseq, next_idx };
-    blockid_t jump_block = { ctx->iseq, jump_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?
     // Maybe we can push the check int into the next block or the stub?
@@ -910,6 +910,12 @@ gen_branchunless(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx)
     x86opnd_t val_opnd = ctx_stack_pop(ctx, 1);
     test(cb, val_opnd, imm_opnd(~Qnil));
 
+    // Get the branch target instruction offsets
+    uint32_t next_idx = ctx_next_idx(ctx);
+    uint32_t jump_idx = next_idx + (uint32_t)ctx_get_arg(ctx, 0);
+    blockid_t next_block = { ctx->iseq, next_idx };
+    blockid_t jump_block = { ctx->iseq, jump_idx };
+
     // Generate the branch instructions
     gen_branch(cb, ocb, jump_block, next_block, gen_branchunless_branch);
 
@@ -947,5 +953,5 @@ ujit_init_codegen(void)
     st_insert(gen_fns, (st_data_t)BIN(opt_minus), (st_data_t)&gen_opt_minus);
     st_insert(gen_fns, (st_data_t)BIN(opt_plus), (st_data_t)&gen_opt_plus);
     st_insert(gen_fns, (st_data_t)BIN(opt_send_without_block), (st_data_t)&gen_opt_send_without_block);
-    st_insert(gen_fns, (st_data_t)BIN(branchunless), (st_data_t)&gen_branchunless);
+    //st_insert(gen_fns, (st_data_t)BIN(branchunless), (st_data_t)&gen_branchunless);
 }
diff --git a/ujit_core.c b/ujit_core.c
index c766f426a4..a1ae75113b 100644
--- a/ujit_core.c
+++ b/ujit_core.c
@@ -163,6 +163,8 @@ uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
     //fprintf(stderr, "rewrite branch at %d\n", branch->start_pos);
 
     // Rewrite the branch with the new jump target address
+    assert (branch->dst_addrs[0] != NULL);
+    assert (branch->dst_addrs[1] != NULL);
     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);