Redo compilation of all ISEQs after invalidation

5 files changed, 58 insertions, 2 deletions

diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 2ad01ef669..5aac3626fa 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -33,6 +33,12 @@ module RubyVM::MJIT
     attr_accessor :write_pos
 
     IseqBlocks = Hash.new { |h, k| h[k] = {} }
+    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))
diff --git a/lib/ruby_vm/mjit/invariants.rb b/lib/ruby_vm/mjit/invariants.rb
index 6dd8e65bc9..d291ade9bf 100644
--- a/lib/ruby_vm/mjit/invariants.rb
+++ b/lib/ruby_vm/mjit/invariants.rb
@@ -58,8 +58,9 @@ module RubyVM::MJIT
       end
 
       def on_tracing_invalidate_all
-        # TODO: assert patches don't overlap each other
+        # On-Stack Replacement
         @patches.each do |address, target|
+          # TODO: assert patches don't overlap each other
           @cb.with_write_addr(address) do
             asm = Assembler.new
             asm.comment('on_tracing_invalidate_all')
@@ -67,6 +68,16 @@ module RubyVM::MJIT
             @cb.write(asm)
           end
         end
+
+        # Avoid reusing past code
+        Compiler.reset_blocks
+
+        C.mjit_for_each_iseq do |iseq|
+          # Disable entering past code
+          iseq.body.jit_func = 0
+          # Compile this again if not converted to trace_* insns
+          iseq.body.total_calls = 0
+        end
       end
 
       private
diff --git a/mjit.c b/mjit.c
index fc19896ff6..97af05d466 100644
--- a/mjit.c
+++ b/mjit.c
@@ -356,7 +356,13 @@ rb_mjit_tracing_invalidate_all(rb_event_flag_t new_iseq_events)
     WITH_MJIT_DISABLED({
         rb_funcall(rb_mMJITHooks, rb_intern("on_tracing_invalidate_all"), 1, UINT2NUM(new_iseq_events));
     });
-    mjit_call_p = false;
+}
+
+// TODO: Use this in more places
+VALUE
+rb_mjit_iseq_new(rb_iseq_t *iseq)
+{
+    return rb_funcall(rb_cMJITIseqPtr, rb_intern("new"), 1, SIZET2NUM((size_t)iseq));
 }
 
 void
diff --git a/mjit_c.c b/mjit_c.c
index 89b0b7a1d0..a97bce05a9 100644
--- a/mjit_c.c
+++ b/mjit_c.c
@@ -15,6 +15,8 @@
 #include "internal.h"
 #include "internal/compile.h"
 #include "internal/hash.h"
+#include "internal/sanitizers.h"
+#include "internal/gc.h"
 #include "yjit.h"
 #include "vm_insnhelper.h"
 
@@ -86,6 +88,33 @@ mjit_enabled_p(rb_execution_context_t *ec, VALUE self)
     return RBOOL(mjit_enabled);
 }
 
+static int
+for_each_iseq_i(void *vstart, void *vend, size_t stride, void *data)
+{
+    VALUE block = (VALUE)data;
+    VALUE v = (VALUE)vstart;
+    for (; v != (VALUE)vend; v += stride) {
+        void *ptr = asan_poisoned_object_p(v);
+        asan_unpoison_object(v, false);
+
+        if (rb_obj_is_iseq(v)) {
+            extern VALUE rb_mjit_iseq_new(rb_iseq_t *iseq);
+            rb_iseq_t *iseq = (rb_iseq_t *)v;
+            rb_funcall(block, rb_intern("call"), 1, rb_mjit_iseq_new(iseq));
+        }
+
+        asan_poison_object_if(ptr, v);
+    }
+    return 0;
+}
+
+static VALUE
+mjit_for_each_iseq(rb_execution_context_t *ec, VALUE self, VALUE block)
+{
+    rb_objspace_each_objects(for_each_iseq_i, (void *)block);
+    return Qnil;
+}
+
 extern bool rb_simple_iseq_p(const rb_iseq_t *iseq);
 
 #include "mjit_c.rbinc"
diff --git a/mjit_c.rb b/mjit_c.rb
index bee063aa3e..6208e4772f 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -145,6 +145,10 @@ module RubyVM::MJIT # :nodoc: all
       }
     end
 
+    def mjit_for_each_iseq(&block)
+      Primitive.mjit_for_each_iseq(block)
+    end
+
     #========================================================================================
     #
     # Old stuff