Step 2 to remove the global cb/ocb objects.

yjit_codegen.c

@@ -596,7 +596,7 @@ jit_jump_to_next_insn(jitstate_t *jit, const ctx_t *current_context)
 
     // Generate the jump instruction
     gen_direct_jump(
-        jit->block,
+        jit,
         &reset_depth,
         jump_block
     );
@@ -630,6 +630,8 @@ yjit_gen_block(block_t *block, rb_execution_context_t *ec)
 
     // Initialize a JIT state object
     jitstate_t jit = {
+        .cb = cb,
+        .ocb = ocb,
         .block = block,
         .iseq = iseq,
         .ec = ec
@@ -1477,7 +1479,7 @@ jit_chain_guard(enum jcc_kinds jcc, jitstate_t *jit, const ctx_t *ctx, uint8_t d
         deeper.chain_depth++;
 
         gen_branch(
-            jit->block,
+            jit,
             ctx,
             (blockid_t) { jit->iseq, jit->insn_idx },
             &deeper,
@@ -1684,7 +1686,7 @@ gen_getinstancevariable(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 {
     // Defer compilation so we can specialize on a runtime `self`
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -1908,7 +1910,7 @@ gen_fixnum_cmp(jitstate_t* jit, ctx_t* ctx, cmov_fn cmov_op)
 {
     // Defer compilation so we can specialize base on a runtime receiver
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -2054,7 +2056,7 @@ gen_opt_eq(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 {
     // Defer compilation so we can specialize base on a runtime receiver
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -2094,7 +2096,7 @@ gen_opt_aref(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 
     // Defer compilation so we can specialize base on a runtime receiver
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -2214,7 +2216,7 @@ gen_opt_aset(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 {
     // Defer compilation so we can specialize on a runtime `self`
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -2293,7 +2295,7 @@ gen_opt_and(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 {
     // Defer compilation so we can specialize on a runtime `self`
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -2336,7 +2338,7 @@ gen_opt_or(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 {
     // Defer compilation so we can specialize on a runtime `self`
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -2379,7 +2381,7 @@ gen_opt_minus(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 {
     // Defer compilation so we can specialize on a runtime `self`
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -2424,7 +2426,7 @@ gen_opt_plus(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 {
     // Defer compilation so we can specialize on a runtime `self`
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -2649,7 +2651,7 @@ gen_branchif(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 
     // Generate the branch instructions
     gen_branch(
-        jit->block,
+        jit,
         ctx,
         jump_block,
         ctx,
@@ -2706,7 +2708,7 @@ gen_branchunless(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 
     // Generate the branch instructions
     gen_branch(
-        jit->block,
+        jit,
         ctx,
         jump_block,
         ctx,
@@ -2762,7 +2764,7 @@ gen_branchnil(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 
     // Generate the branch instructions
     gen_branch(
-        jit->block,
+        jit,
         ctx,
         jump_block,
         ctx,
@@ -2791,7 +2793,7 @@ gen_jump(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 
     // Generate the jump instruction
     gen_direct_jump(
-        jit->block,
+        jit,
         ctx,
         jump_block
     );
@@ -3516,7 +3518,7 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
 
     // Write the JIT return address on the callee frame
     gen_branch(
-        jit->block,
+        jit,
         ctx,
         return_block,
         &return_ctx,
@@ -3533,7 +3535,7 @@ gen_send_iseq(jitstate_t *jit, ctx_t *ctx, const struct rb_callinfo *ci, const r
 
     // Directly jump to the entry point of the callee
     gen_direct_jump(
-        jit->block,
+        jit,
         &callee_ctx,
         (blockid_t){ iseq, start_pc_offset }
     );
@@ -3577,7 +3579,7 @@ gen_send_general(jitstate_t *jit, ctx_t *ctx, struct rb_call_data *cd, rb_iseq_t
 
     // Defer compilation so we can specialize on class of receiver
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -3710,7 +3712,7 @@ gen_invokesuper(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
 
     // Defer compilation so we can specialize on class of receiver
     if (!jit_at_current_insn(jit)) {
-        defer_compilation(jit->block, jit->insn_idx, ctx);
+        defer_compilation(jit, ctx);
         return YJIT_END_BLOCK;
     }
 
@@ -4068,7 +4070,7 @@ gen_opt_getinlinecache(jitstate_t* jit, ctx_t* ctx, codeblock_t* cb)
     // Jump over the code for filling the cache
     uint32_t jump_idx = jit_next_insn_idx(jit) + (int32_t)jump_offset;
     gen_direct_jump(
-        jit->block,
+        jit,
         ctx,
         (blockid_t){ .iseq = jit->iseq, .idx = jump_idx }
     );

yjit_codegen.h

@@ -9,37 +9,6 @@ extern codeblock_t *cb;
 extern codeblock_t *ocb;
 extern uint32_t yjit_codepage_frozen_bytes;
 
-// Code generation state
-typedef struct JITState
-{
-    // Block version being compiled
-    block_t *block;
-
-    // Instruction sequence this is associated with
-    const rb_iseq_t *iseq;
-
-    // Index of the current instruction being compiled
-    uint32_t insn_idx;
-
-    // Opcode for the instruction being compiled
-    int opcode;
-
-    // PC of the instruction being compiled
-    VALUE *pc;
-
-    // Side exit to the instruction being compiled. See :side-exit:.
-    uint8_t *side_exit_for_pc;
-
-    // Execution context when compilation started
-    // This allows us to peek at run-time values
-    rb_execution_context_t *ec;
-
-    // Whether we need to record the code address at
-    // the end of this bytecode instruction for global invalidation
-    bool record_boundary_patch_point;
-
-} jitstate_t;
-
 typedef enum codegen_status {
     YJIT_END_BLOCK,
     YJIT_KEEP_COMPILING,

yjit_core.c

@@ -871,7 +871,7 @@ uint8_t* get_branch_target(
 }
 
 void gen_branch(
-    block_t* block,
+    jitstate_t* jit,
     const ctx_t* src_ctx,
     blockid_t target0,
     const ctx_t* ctx0,
@@ -882,7 +882,7 @@ void gen_branch(
 {
     RUBY_ASSERT(target0.iseq != NULL);
 
-    branch_t* branch = make_branch_entry(block, src_ctx, gen_fn);
+    branch_t* branch = make_branch_entry(jit->block, src_ctx, gen_fn);
     branch->targets[0] = target0;
     branch->targets[1] = target1;
     branch->target_ctxs[0] = *ctx0;
@@ -917,14 +917,14 @@ gen_jump_branch(codeblock_t* cb, uint8_t* target0, uint8_t* target1, uint8_t sha
 }
 
 void gen_direct_jump(
-    block_t* block,
+    jitstate_t* jit,
     const ctx_t* ctx,
     blockid_t target0
 )
 {
     RUBY_ASSERT(target0.iseq != NULL);
 
-    branch_t* branch = make_branch_entry(block, ctx, gen_jump_branch);
+    branch_t* branch = make_branch_entry(jit->block, ctx, gen_jump_branch);
     branch->targets[0] = target0;
     branch->target_ctxs[0] = *ctx;
 
@@ -955,8 +955,7 @@ void gen_direct_jump(
 
 // Create a stub to force the code up to this point to be executed
 void defer_compilation(
-    block_t* block,
+    jitstate_t* jit,
-    uint32_t insn_idx,
     ctx_t* cur_ctx
 )
 {
@@ -974,14 +973,15 @@ void defer_compilation(
 
     next_ctx.chain_depth += 1;
 
-    branch_t* branch = make_branch_entry(block, cur_ctx, gen_jump_branch);
+    branch_t* branch = make_branch_entry(jit->block, cur_ctx, gen_jump_branch);
 
     // Get the branch targets or stubs
     branch->target_ctxs[0] = next_ctx;
-    branch->targets[0] = (blockid_t){ block->blockid.iseq, insn_idx };
+    branch->targets[0] = (blockid_t){ jit->block->blockid.iseq, jit->insn_idx };
     branch->dst_addrs[0] = get_branch_target(branch->targets[0], &next_ctx, branch, 0);
 
     // Call the branch generation function
+    codeblock_t* cb = jit->cb;
     branch->start_pos = cb->write_pos;
     gen_jump_branch(cb, branch->dst_addrs[0], NULL, SHAPE_DEFAULT);
     branch->end_pos = cb->write_pos;

yjit_core.h

@@ -263,6 +263,42 @@ typedef struct yjit_block_version
 
 } block_t;
 
+// Code generation state
+typedef struct JITState
+{
+    // Inline and outlined code blocks we are
+    // currently generating code into
+    codeblock_t* cb;
+    codeblock_t* ocb;
+
+    // Block version being compiled
+    block_t *block;
+
+    // Instruction sequence this is associated with
+    const rb_iseq_t *iseq;
+
+    // Index of the current instruction being compiled
+    uint32_t insn_idx;
+
+    // Opcode for the instruction being compiled
+    int opcode;
+
+    // PC of the instruction being compiled
+    VALUE *pc;
+
+    // Side exit to the instruction being compiled. See :side-exit:.
+    uint8_t *side_exit_for_pc;
+
+    // Execution context when compilation started
+    // This allows us to peek at run-time values
+    rb_execution_context_t *ec;
+
+    // Whether we need to record the code address at
+    // the end of this bytecode instruction for global invalidation
+    bool record_boundary_patch_point;
+
+} jitstate_t;
+
 // Context object methods
 x86opnd_t ctx_sp_opnd(ctx_t* ctx, int32_t offset_bytes);
 x86opnd_t ctx_stack_push_mapping(ctx_t* ctx, temp_type_mapping_t mapping);
@@ -290,7 +326,7 @@ void yjit_free_block(block_t *block);
 rb_yjit_block_array_t yjit_get_version_array(const rb_iseq_t *iseq, unsigned idx);
 
 void gen_branch(
-    block_t* block,
+    jitstate_t* jit,
     const ctx_t* src_ctx,
     blockid_t target0,
     const ctx_t* ctx0,
@@ -300,14 +336,13 @@ void gen_branch(
 );
 
 void gen_direct_jump(
-    block_t* block,
+    jitstate_t* jit,
     const ctx_t* ctx,
     blockid_t target0
 );
 
 void defer_compilation(
-    block_t* block,
+    jitstate_t* jit,
-    uint32_t insn_idx,
     ctx_t* cur_ctx
 );