Use a versioning context when compiling blocks

4 files changed, 18 insertions, 12 deletions

diff --git a/ujit_codegen.c b/ujit_codegen.c
index a1fcb8b7c0..7c5a2bd72b 100644
--- a/ujit_codegen.c
+++ b/ujit_codegen.c
@@ -129,9 +129,12 @@ uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx)
     // Write the interpreter entry prologue
     ujit_gen_entry(cb);
 
+    // Create codegen context
+    ctx_t ctx = { 0 };
+
     // Compile the block starting at this instruction
     uint32_t num_instrs = 0;
-    ujit_compile_block(iseq, insn_idx, &num_instrs);
+    ujit_compile_block(iseq, insn_idx, &ctx, &num_instrs);
 
     // If no instructions were compiled
     if (num_instrs == 0) {
@@ -152,7 +155,7 @@ uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx)
 Compile a sequence of bytecode instructions starting at `insn_idx`.
 */
 uint8_t *
-ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, uint32_t* num_instrs)
+ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx, uint32_t* num_instrs)
 {
     assert (cb != NULL);
     VALUE *encoded = iseq->body->iseq_encoded;
@@ -175,9 +178,6 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, uint32_t* num_instr
     jit.iseq = iseq;
     jit.start_idx = insn_idx;
 
-    // Create codegen context
-    ctx_t ctx = { 0 };
-
     // For each instruction to compile
     for (;;) {
         // Set the current instruction
@@ -201,7 +201,7 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, uint32_t* num_instr
 
         // Call the code generation function
         codegen_fn gen_fn = (codegen_fn)st_gen_fn;
-        if (!gen_fn(&jit, &ctx)) {
+        if (!gen_fn(&jit, ctx)) {
             break;
         }
 
@@ -222,7 +222,7 @@ ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, uint32_t* num_instr
     // FIXME: we only need to generate an exit if an instruction fails to compile
     //
     // Generate code to exit to the interpreter
-    ujit_gen_exit(&jit, &ctx, cb, &encoded[insn_idx]);
+    ujit_gen_exit(&jit, ctx, cb, &encoded[insn_idx]);
 
     if (UJIT_DUMP_MODE >= 2) {
         // Dump list of compiled instrutions
@@ -958,7 +958,7 @@ gen_branchunless(jitstate_t* jit, ctx_t* ctx)
     blockid_t jump_block = { jit->iseq, jump_idx };
 
     // Generate the branch instructions
-    gen_branch(cb, ocb, jump_block, next_block, gen_branchunless_branch);
+    gen_branch(ctx, jump_block, next_block, gen_branchunless_branch);
 
     return true;
 }
diff --git a/ujit_codegen.h b/ujit_codegen.h
index 3d0d9d7e03..7abeaadfc0 100644
--- a/ujit_codegen.h
+++ b/ujit_codegen.h
@@ -2,6 +2,7 @@
 #define UJIT_CODEGEN_H 1
 
 #include "stddef.h"
+#include "ujit_core.h"
 
 // Code blocks we generate code into
 codeblock_t* cb;
@@ -29,7 +30,7 @@ typedef bool (*codegen_fn)(jitstate_t* jit, ctx_t* ctx);
 
 uint8_t* ujit_compile_entry(const rb_iseq_t *iseq, uint32_t insn_idx);
 
-uint8_t *ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, uint32_t* num_instrs);
+uint8_t *ujit_compile_block(const rb_iseq_t *iseq, uint32_t insn_idx, ctx_t* ctx, uint32_t* num_instrs);
 
 void ujit_init_codegen(void);
 
diff --git a/ujit_core.c b/ujit_core.c
index 0bffec7ebf..af0b8c86a5 100644
--- a/ujit_core.c
+++ b/ujit_core.c
@@ -133,8 +133,9 @@ uint8_t* branch_stub_hit(uint32_t branch_idx, uint32_t target_idx)
     {
         //fprintf(stderr, "compiling block\n");
 
+        ctx_t ctx = branch->ctx;
         uint32_t num_instrs = 0;
-        block_ptr = ujit_compile_block(target.iseq, target.idx, &num_instrs);
+        block_ptr = ujit_compile_block(target.iseq, target.idx, &ctx, &num_instrs);
         st_insert(version_tbl, (st_data_t)&target, (st_data_t)block_ptr);
         branch->dst_addrs[target_idx] = block_ptr;
     }
@@ -192,7 +193,7 @@ uint8_t* get_branch_target(codeblock_t* ocb, blockid_t target, uint32_t branch_i
     return stub_addr;
 }
 
-void gen_branch(codeblock_t* cb, codeblock_t* ocb, blockid_t target0, blockid_t target1, branchgen_fn gen_fn)
+void gen_branch(ctx_t* ctx, blockid_t target0, blockid_t target1, branchgen_fn gen_fn)
 {
     // Get branch targets or stubs (code pointers)
     uint8_t* dst_addr0 = get_branch_target(ocb, target0, num_branches, 0);
@@ -207,6 +208,7 @@ void gen_branch(codeblock_t* cb, codeblock_t* ocb, blockid_t target0, blockid_t
 
     // Register this branch entry
     branch_t branch_entry = {
+        *ctx,
         start_pos,
         end_pos,
         { target0, target1 },
diff --git a/ujit_core.h b/ujit_core.h
index dd2ea63d40..d508017649 100644
--- a/ujit_core.h
+++ b/ujit_core.h
@@ -56,6 +56,9 @@ typedef void (*branchgen_fn)(codeblock_t* cb, uint8_t* target0, uint8_t* target1
 // Store info about an outgoing branch in a code segment
 typedef struct BranchEntry
 {
+    // Context right after the branch instruction
+    ctx_t ctx;
+
     // Positions where the generated code starts and ends
     uint32_t start_pos;
     uint32_t end_pos;
@@ -83,7 +86,7 @@ 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);
 
-void gen_branch(codeblock_t* cb, codeblock_t* ocb, blockid_t target0, blockid_t target1, branchgen_fn gen_fn);
+void gen_branch(ctx_t* ctx, blockid_t target0, blockid_t target1, branchgen_fn gen_fn);
 
 void ujit_init_core(void);