ZJIT: Panic on BOP redefinition only when needed (#13782)

.github/workflows/zjit-macos.yml

@@ -119,6 +119,7 @@ jobs:
           ../src/bootstraptest/test_flow.rb \
           ../src/bootstraptest/test_fork.rb \
           ../src/bootstraptest/test_gc.rb \
+          ../src/bootstraptest/test_insns.rb \
           ../src/bootstraptest/test_io.rb \
           ../src/bootstraptest/test_jump.rb \
           ../src/bootstraptest/test_literal.rb \
@@ -138,7 +139,6 @@ jobs:
           ../src/bootstraptest/test_yjit_30k_methods.rb \
           ../src/bootstraptest/test_yjit_rust_port.rb
         # ../src/bootstraptest/test_eval.rb \
-        # ../src/bootstraptest/test_insns.rb \
         # ../src/bootstraptest/test_yjit.rb \
         if: ${{ matrix.test_task == 'btest' }}
 

.github/workflows/zjit-ubuntu.yml

@@ -141,6 +141,7 @@ jobs:
           ../src/bootstraptest/test_flow.rb \
           ../src/bootstraptest/test_fork.rb \
           ../src/bootstraptest/test_gc.rb \
+          ../src/bootstraptest/test_insns.rb \
           ../src/bootstraptest/test_io.rb \
           ../src/bootstraptest/test_jump.rb \
           ../src/bootstraptest/test_literal.rb \
@@ -160,7 +161,6 @@ jobs:
           ../src/bootstraptest/test_yjit_30k_methods.rb \
           ../src/bootstraptest/test_yjit_rust_port.rb
         # ../src/bootstraptest/test_eval.rb \
-        # ../src/bootstraptest/test_insns.rb \
         # ../src/bootstraptest/test_yjit.rb \
         if: ${{ matrix.test_task == 'btest' }}
 

test/ruby/test_zjit.rb

@@ -815,6 +815,21 @@ class TestZJIT < Test::Unit::TestCase
     }, call_threshold: 2
   end
 
+  def test_bop_invalidation
+    omit 'Invalidation on BOP redefinition is not implemented yet'
+    assert_compiles '', %q{
+      def test
+        eval(<<~RUBY)
+          class Integer
+            def +(_) = 100
+          end
+        RUBY
+        1 + 2
+      end
+      test
+    }
+  end
+
   def test_defined_yield
     assert_compiles "nil", "defined?(yield)"
     assert_compiles '[nil, nil, "yield"]', %q{

zjit/src/backend/arm64/mod.rs

@@ -892,7 +892,6 @@ impl Assembler
         let mut pos_markers: Vec<(usize, CodePtr)> = vec![];
 
         // For each instruction
-        //let start_write_pos = cb.get_write_pos();
         let mut insn_idx: usize = 0;
         while let Some(insn) = self.insns.get(insn_idx) {
             //let src_ptr = cb.get_write_ptr();
@@ -1256,9 +1255,6 @@ impl Assembler
                     csel(cb, out.into(), truthy.into(), falsy.into(), Condition::GE);
                 }
                 Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
-                Insn::PadInvalPatch => {
-                    unimplemented!("we haven't needed padding in ZJIT yet");
-                }
             };
 
             // On failure, jump to the next page and retry the current insn

zjit/src/backend/lir.rs

@@ -484,10 +484,6 @@ pub enum Insn {
     // binary OR operation.
     Or { left: Opnd, right: Opnd, out: Opnd },
 
-    /// Pad nop instructions to accommodate Op::Jmp in case the block or the insn
-    /// is invalidated.
-    PadInvalPatch,
-
     // Mark a position in the generated code
     PosMarker(PosMarkerFn),
 
@@ -607,7 +603,6 @@ impl Insn {
             Insn::Mov { .. } => "Mov",
             Insn::Not { .. } => "Not",
             Insn::Or { .. } => "Or",
-            Insn::PadInvalPatch => "PadEntryExit",
             Insn::PosMarker(_) => "PosMarker",
             Insn::RShift { .. } => "RShift",
             Insn::Store { .. } => "Store",
@@ -801,7 +796,6 @@ impl<'a> Iterator for InsnOpndIterator<'a> {
             Insn::CPushAll |
             Insn::FrameSetup |
             Insn::FrameTeardown |
-            Insn::PadInvalPatch |
             Insn::PosMarker(_) => None,
 
             Insn::CPopInto(opnd) |
@@ -956,7 +950,6 @@ impl<'a> InsnOpndMutIterator<'a> {
             Insn::CPushAll |
             Insn::FrameSetup |
             Insn::FrameTeardown |
-            Insn::PadInvalPatch |
             Insn::PosMarker(_) => None,
 
             Insn::CPopInto(opnd) |
@@ -2171,10 +2164,6 @@ impl Assembler {
         out
     }
 
-    pub fn pad_inval_patch(&mut self) {
-        self.push_insn(Insn::PadInvalPatch);
-    }
-
     //pub fn pos_marker<F: FnMut(CodePtr)>(&mut self, marker_fn: F)
     pub fn pos_marker(&mut self, marker_fn: impl Fn(CodePtr, &CodeBlock) + 'static) {
         self.push_insn(Insn::PosMarker(Box::new(marker_fn)));

zjit/src/backend/x86_64/mod.rs

@@ -444,7 +444,6 @@ impl Assembler
         let mut pos_markers: Vec<(usize, CodePtr)> = vec![];
 
         // For each instruction
-        //let start_write_pos = cb.get_write_pos();
         let mut insn_idx: usize = 0;
         while let Some(insn) = self.insns.get(insn_idx) {
             //let src_ptr = cb.get_write_ptr();
@@ -795,15 +794,6 @@ impl Assembler
                     emit_csel(cb, *truthy, *falsy, *out, cmovge, cmovl);
                 }
                 Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
-                Insn::PadInvalPatch => {
-                    unimplemented!("we don't need padding yet");
-                    /*
-                    let code_size = cb.get_write_pos().saturating_sub(std::cmp::max(start_write_pos, cb.page_start_pos()));
-                    if code_size < cb.jmp_ptr_bytes() {
-                        nop(cb, (cb.jmp_ptr_bytes() - code_size) as u32);
-                    }
-                    */
-                }
             };
 
             // On failure, jump to the next page and retry the current insn

zjit/src/codegen.rs

@@ -3,11 +3,12 @@ use std::rc::Rc;
 use std::num::NonZeroU32;
 
 use crate::backend::current::{Reg, ALLOC_REGS};
+use crate::invariants::track_bop_assumption;
 use crate::profile::get_or_create_iseq_payload;
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, SP};
-use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType, SELF_PARAM_IDX, SpecialObjectType};
+use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, Invariant, RangeType, SpecialObjectType, SELF_PARAM_IDX};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
 use crate::hir_type::{types::Fixnum, Type};
 use crate::options::get_option;
@@ -286,7 +287,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::Test { val } => gen_test(asm, opnd!(val))?,
         Insn::GuardType { val, guard_type, state } => gen_guard_type(jit, asm, opnd!(val), *guard_type, &function.frame_state(*state))?,
         Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(jit, asm, opnd!(val), *expected, &function.frame_state(*state))?,
-        Insn::PatchPoint(_) => return Some(()), // For now, rb_zjit_bop_redefined() panics. TODO: leave a patch point and fix rb_zjit_bop_redefined()
+        Insn::PatchPoint(invariant) => return gen_patch_point(asm, invariant),
         Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(jit, asm, *cfun, args)?,
         Insn::GetIvar { self_val, id, state: _ } => gen_getivar(asm, opnd!(self_val), *id),
         Insn::SetGlobal { id, val, state: _ } => return Some(gen_setglobal(asm, *id, opnd!(val))),
@@ -422,6 +423,24 @@ fn gen_invokebuiltin(jit: &mut JITState, asm: &mut Assembler, state: &FrameState
     Some(val)
 }
 
+/// Record a patch point that should be invalidated on a given invariant
+fn gen_patch_point(asm: &mut Assembler, invariant: &Invariant) -> Option<()> {
+    let invariant = invariant.clone();
+    asm.pos_marker(move |code_ptr, _cb| {
+        match invariant {
+            Invariant::BOPRedefined { klass, bop } => {
+                track_bop_assumption(klass, bop, code_ptr);
+            }
+            _ => {
+                debug!("ZJIT: gen_patch_point: unimplemented invariant {invariant:?}");
+                return;
+            }
+        }
+    });
+    // TODO: Make sure patch points do not overlap with each other.
+    Some(())
+}
+
 /// Lowering for [`Insn::CCall`]. This is a low-level raw call that doesn't know
 /// anything about the callee, so handling for e.g. GC safety is dealt with elsewhere.
 fn gen_ccall(jit: &mut JITState, asm: &mut Assembler, cfun: *const u8, args: &[InsnId]) -> Option<lir::Opnd> {

zjit/src/invariants.rs

@@ -1,6 +1,6 @@
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
-use crate::{cruby::{ruby_basic_operators, IseqPtr, RedefinitionFlag}, state::ZJITState, state::zjit_enabled_p};
+use crate::{cruby::{ruby_basic_operators, IseqPtr, RedefinitionFlag}, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
 
 /// Used to track all of the various block references that contain assumptions
 /// about the state of the virtual machine.
@@ -11,19 +11,28 @@ pub struct Invariants {
 
     /// Set of ISEQs whose JIT code assumes that it doesn't escape EP
     no_ep_escape_iseqs: HashSet<IseqPtr>,
+
+    /// Map from a class and its associated basic operator to a set of patch points
+    bop_patch_points: HashMap<(RedefinitionFlag, ruby_basic_operators), HashSet<CodePtr>>,
 }
 
 /// Called when a basic operator is redefined. Note that all the blocks assuming
 /// the stability of different operators are invalidated together and we don't
 /// do fine-grained tracking.
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_bop_redefined(_klass: RedefinitionFlag, _bop: ruby_basic_operators) {
+pub extern "C" fn rb_zjit_bop_redefined(klass: RedefinitionFlag, bop: ruby_basic_operators) {
     // If ZJIT isn't enabled, do nothing
     if !zjit_enabled_p() {
         return;
     }
 
-    unimplemented!("Invalidation on BOP redefinition is not implemented yet");
+    let invariants = ZJITState::get_invariants();
+    if let Some(code_ptrs) = invariants.bop_patch_points.get(&(klass, bop)) {
+        // Invalidate all patch points for this BOP
+        for &ptr in code_ptrs {
+            unimplemented!("Invalidation on BOP redefinition is not implemented yet: {ptr:?}");
+        }
+    }
 }
 
 /// Invalidate blocks for a given ISEQ that assumes environment pointer is
@@ -57,3 +66,9 @@ pub fn track_no_ep_escape_assumption(iseq: IseqPtr) {
 pub fn iseq_escapes_ep(iseq: IseqPtr) -> bool {
     ZJITState::get_invariants().ep_escape_iseqs.contains(&iseq)
 }
+
+/// Track a patch point for a basic operator in a given class.
+pub fn track_bop_assumption(klass: RedefinitionFlag, bop: ruby_basic_operators, code_ptr: CodePtr) {
+    let invariants = ZJITState::get_invariants();
+    invariants.bop_patch_points.entry((klass, bop)).or_default().insert(code_ptr);
+}

zjit/src/virtualmem.rs

@@ -62,7 +62,7 @@ pub trait Allocator {
 
 /// Pointer into a [VirtualMemory] represented as an offset from the base.
 /// Note: there is no NULL constant for [CodePtr]. You should use `Option<CodePtr>` instead.
-#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Debug)]
+#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Debug)]
 #[repr(C, packed)]
 pub struct CodePtr(u32);