1 files changed, 27 insertions, 1 deletions
diff --git a/lib/ruby_vm/mjit/invariants.rb b/lib/ruby_vm/mjit/invariants.rb
index c9d9b9fb6a..fdd39ba466 100644
--- a/lib/ruby_vm/mjit/invariants.rb
+++ b/lib/ruby_vm/mjit/invariants.rb
@@ -2,12 +2,22 @@ require 'set'
 
 module RubyVM::MJIT
   class Invariants
+    # @param cb [CodeBlock]
     # @param ocb [CodeBlock]
     # @param exit_compiler [RubyVM::MJIT::ExitCompiler]
-    def initialize(ocb, exit_compiler)
+    def initialize(cb, ocb, exit_compiler)
+      @cb = cb
       @ocb = ocb
       @exit_compiler = exit_compiler
       @bop_blocks = Set.new # TODO: actually invalidate this
+      @cme_blocks = Hash.new { |h, k| h[k] = Set.new }
+
+      invariants = self
+      hooks = Module.new
+      hooks.define_method(:on_cme_invalidate) do |cme|
+        invariants.on_cme_invalidate(cme)
+      end
+      Hooks.singleton_class.prepend(hooks)
     end
 
     # @param jit [RubyVM::MJIT::JITState]
@@ -21,6 +31,22 @@ module RubyVM::MJIT
       true
     end
 
+    # @param jit [RubyVM::MJIT::JITState]
+    def assume_method_lookup_stable(jit, cme)
+      ensure_block_entry_exit(jit.block, cause: 'assume_method_lookup_stable')
+      @cme_blocks[cme.to_i] << jit.block
+    end
+
+    def on_cme_invalidate(cme)
+      @cme_blocks.fetch(cme.to_i, []).each do |block|
+        @cb.with_write_addr(block.start_addr) do
+          asm = Assembler.new
+          asm.jmp(block.entry_exit)
+          @cb.write(asm)
+        end
+      end
+    end
+
     private
 
     # @param block [RubyVM::MJIT::Block]