Store MJIT blocks on each ISEQ

Merged: https://github.com/ruby/ruby/pull/7448

iseq.c

@@ -357,7 +357,7 @@ rb_iseq_mark_and_move(rb_iseq_t *iseq, bool reference_updating)
 
         if (reference_updating) {
 #if USE_MJIT
-            rb_mjit_iseq_update_references(iseq);
+            rb_mjit_iseq_update_references(body);
 #endif
 #if USE_YJIT
             rb_yjit_iseq_update_references(body->yjit_payload);
@@ -365,7 +365,7 @@ rb_iseq_mark_and_move(rb_iseq_t *iseq, bool reference_updating)
         }
         else {
 #if USE_MJIT
-            mjit_mark_cc_entries(body);
+            rb_mjit_iseq_mark(body->mjit_blocks);
 #endif
 #if USE_YJIT
             rb_yjit_iseq_mark(body->yjit_payload);

lib/ruby_vm/mjit/compiler.rb

@@ -29,21 +29,12 @@ module RubyVM::MJIT
 
   # Scratch registers: rax, rcx
 
+  # Mark objects in this Array during GC
+  GC_REFS = []
+
   class Compiler
     attr_accessor :write_pos
 
-    IseqBlocks = Hash.new do |iseq_pc_ctx, iseq|
-      iseq_pc_ctx[iseq] = Hash.new do |pc_ctx, pc|
-        pc_ctx[pc] = {}
-      end
-    end
-    DeadBlocks = [] # invalidated IseqBlocks, but kept for safety
-
-    def self.reset_blocks
-      DeadBlocks << IseqBlocks.dup
-      IseqBlocks.clear
-    end
-
     def self.decode_insn(encoded)
       INSNS.fetch(C.rb_vm_insn_decode(encoded))
     end
@@ -268,24 +259,34 @@ module RubyVM::MJIT
     end
 
     def list_blocks(iseq, pc)
-      IseqBlocks[iseq.to_i][pc].values
+      mjit_blocks(iseq)[pc].values
     end
 
     # @param [Integer] pc
     # @param [RubyVM::MJIT::Context] ctx
     # @return [RubyVM::MJIT::Block,NilClass]
     def find_block(iseq, pc, ctx)
-      IseqBlocks[iseq.to_i][pc][ctx]
+      mjit_blocks(iseq)[pc][ctx]
     end
 
     # @param [RubyVM::MJIT::Block] block
     def set_block(iseq, block)
-      IseqBlocks[iseq.to_i][block.pc][block.ctx] = block
+      mjit_blocks(iseq)[block.pc][block.ctx] = block
     end
 
     # @param [RubyVM::MJIT::Block] block
     def remove_block(iseq, block)
-      IseqBlocks[iseq.to_i][block.pc].delete(block.ctx)
+      mjit_blocks(iseq)[block.pc].delete(block.ctx)
+    end
+
+    def mjit_blocks(iseq)
+      unless iseq.body.mjit_blocks
+        iseq.body.mjit_blocks = Hash.new { |h, k| h[k] = {} }
+        # For some reason, rb_mjit_iseq_mark didn't protect this Hash
+        # from being freed. So we rely on GC_REFS to keep the Hash.
+        GC_REFS << iseq.body.mjit_blocks
+      end
+      iseq.body.mjit_blocks
     end
   end
 end

lib/ruby_vm/mjit/invariants.rb

@@ -112,12 +112,11 @@ module RubyVM::MJIT
         end
         @patches.clear
 
-        # Avoid reusing past code
-        Compiler.reset_blocks
-
         C.mjit_for_each_iseq do |iseq|
-          # Disable entering past code
+          # Avoid entering past code
           iseq.body.jit_func = 0
+          # Avoid reusing past code
+          iseq.body.mjit_blocks.clear if iseq.body.mjit_blocks
           # Compile this again if not converted to trace_* insns
           iseq.body.total_calls = 0
         end

mjit.c

@@ -272,12 +272,6 @@ mjit_child_after_fork(void)
     // TODO: remove this
 }
 
-void
-mjit_mark_cc_entries(const struct rb_iseq_constant_body *const body)
-{
-    // TODO: implement
-}
-
 // Compile ISeq to C code in `f`. It returns true if it succeeds to compile.
 bool
 mjit_compile(FILE *f, const rb_iseq_t *iseq, const char *funcname, int id)
@@ -394,11 +388,13 @@ mjit_iseq_update_references(void *data)
 }
 
 void
-rb_mjit_iseq_update_references(const rb_iseq_t *iseq)
+rb_mjit_iseq_update_references(struct rb_iseq_constant_body *const body)
 {
     if (!mjit_enabled) return;
 
-    // TODO: update mjit_blocks
+    if (body->mjit_blocks) {
+        body->mjit_blocks = rb_gc_location(body->mjit_blocks);
+    }
 
     // Asynchronously hook the Ruby code to avoid allocation during GC.compact.
     // Using _one because it's too slow to invalidate all for each ISEQ. Thus
@@ -406,6 +402,18 @@ rb_mjit_iseq_update_references(const rb_iseq_t *iseq)
     rb_postponed_job_register_one(0, mjit_iseq_update_references, NULL);
 }
 
+void
+rb_mjit_iseq_mark(VALUE mjit_blocks)
+{
+    if (!mjit_enabled) return;
+
+    // Note: This wasn't enough for some reason.
+    // We actually rely on RubyVM::MJIT::GC_REFS to mark this.
+    if (mjit_blocks) {
+        rb_gc_mark_movable(mjit_blocks);
+    }
+}
+
 // TODO: Use this in more places
 VALUE
 rb_mjit_iseq_new(rb_iseq_t *iseq)

mjit.h

@@ -106,9 +106,9 @@ extern void mjit_cancel_all(const char *reason);
 extern bool mjit_compile(FILE *f, const rb_iseq_t *iseq, const char *funcname, int id);
 extern void mjit_init(const struct mjit_options *opts);
 extern void mjit_free_iseq(const rb_iseq_t *iseq);
-extern void rb_mjit_iseq_update_references(const rb_iseq_t *iseq);
+extern void rb_mjit_iseq_update_references(struct rb_iseq_constant_body *const body);
 extern void mjit_mark(void);
-extern void mjit_mark_cc_entries(const struct rb_iseq_constant_body *const body);
+extern void rb_mjit_iseq_mark(VALUE mjit_blocks);
 extern void mjit_notify_waitpid(int exit_code);
 
 extern void rb_mjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);

vm_core.h

@@ -511,7 +511,7 @@ struct rb_iseq_constant_body {
 #endif
 
 #if USE_MJIT
-    // MJIT stores { Context => Block } for each iseq.
+    // MJIT stores some data on each iseq.
     VALUE mjit_blocks;
 #endif