ZJIT: Replace GetConstantPath with Const if the IC is not empty (#13183)

* Add rb_zjit_constcache_shareable

* Add rb_zjit_multi_ractor_p

* Replace GetConstantPath with Const if the IC is not empty

zjit.c

@@ -525,6 +525,12 @@ rb_BASIC_OP_UNREDEFINED_P(enum ruby_basic_operators bop, uint32_t klass)
     return BASIC_OP_UNREDEFINED_P(bop, klass);
 }
 
+bool
+rb_zjit_multi_ractor_p(void)
+{
+    return rb_multi_ractor_p();
+}
+
 // For debug builds
 void
 rb_assert_iseq_handle(VALUE handle)
@@ -532,6 +538,12 @@ rb_assert_iseq_handle(VALUE handle)
     RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(handle, imemo_iseq));
 }
 
+bool
+rb_zjit_constcache_shareable(const struct iseq_inline_constant_cache_entry *ice)
+{
+    return (ice->flags & IMEMO_CONST_CACHE_SHAREABLE) != 0;
+}
+
 void
 rb_assert_cme_handle(VALUE handle)
 {

zjit/bindgen/src/main.rs

@@ -409,6 +409,8 @@ fn main() {
         .allowlist_function("rb_get_cfp_ep")
         .allowlist_function("rb_get_cfp_ep_level")
         .allowlist_function("rb_get_cme_def_type")
+        .allowlist_function("rb_zjit_multi_ractor_p")
+        .allowlist_function("rb_zjit_constcache_shareable")
         .allowlist_function("rb_get_cme_def_body_attr_id")
         .allowlist_function("rb_get_symbol_id")
         .allowlist_function("rb_get_cme_def_body_optimized_type")

zjit/src/cruby_bindings.inc.rs

@@ -995,7 +995,9 @@ unsafe extern "C" {
     pub fn rb_RB_TYPE_P(obj: VALUE, t: ruby_value_type) -> bool;
     pub fn rb_RSTRUCT_LEN(st: VALUE) -> ::std::os::raw::c_long;
     pub fn rb_BASIC_OP_UNREDEFINED_P(bop: ruby_basic_operators, klass: u32) -> bool;
+    pub fn rb_zjit_multi_ractor_p() -> bool;
     pub fn rb_assert_iseq_handle(handle: VALUE);
+    pub fn rb_zjit_constcache_shareable(ice: *const iseq_inline_constant_cache_entry) -> bool;
     pub fn rb_assert_cme_handle(handle: VALUE);
     pub fn rb_IMEMO_TYPE_P(imemo: VALUE, imemo_type: imemo_type) -> ::std::os::raw::c_int;
     pub fn rb_zjit_vm_unlock(

zjit/src/hir.rs

@@ -113,6 +113,13 @@ pub enum Invariant {
         /// The method ID of the method we want to assume unchanged
         method: ID,
     },
+    /// A list of constant expression path segments that must have not been written to for the
+    /// following code to be valid.
+    StableConstantNames {
+        idlist: *const ID,
+    },
+    /// There is one ractor running. If a non-root ractor gets spawned, this is invalidated.
+    SingleRactorMode,
 }
 
 impl Invariant {
@@ -161,6 +168,22 @@ impl<'a> std::fmt::Display for InvariantPrinter<'a> {
                     self.ptr_map.map_id(method.0)
                 )
             }
+            Invariant::StableConstantNames { idlist } => {
+                write!(f, "StableConstantNames({:p}, ", self.ptr_map.map_ptr(idlist))?;
+                let mut idx = 0;
+                let mut sep = "";
+                loop {
+                    let id = unsafe { *idlist.wrapping_add(idx) };
+                    if id.0 == 0 {
+                        break;
+                    }
+                    write!(f, "{sep}{}", id.contents_lossy())?;
+                    sep = "::";
+                    idx += 1;
+                }
+                write!(f, ")")
+            }
+            Invariant::SingleRactorMode => write!(f, "SingleRactorMode"),
         }
     }
 }
@@ -292,7 +315,7 @@ pub enum Insn {
     // with IfTrue/IfFalse in the backend to generate jcc.
     Test { val: InsnId },
     Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId },
-    GetConstantPath { ic: *const u8 },
+    GetConstantPath { ic: *const iseq_inline_constant_cache },
 
     //NewObject?
     //SetIvar {},
@@ -1013,6 +1036,25 @@ impl Function {
                         let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { self_val, call_info, cd, iseq, args, state });
                         self.make_equal_to(insn_id, send_direct);
                     }
+                    Insn::GetConstantPath { ic } => {
+                        let idlist: *const ID = unsafe { (*ic).segments };
+                        let ice = unsafe { (*ic).entry };
+                        if ice.is_null() {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        let cref_sensitive = !unsafe { (*ice).ic_cref }.is_null();
+                        let multi_ractor_mode = unsafe { rb_zjit_multi_ractor_p() };
+                        if cref_sensitive || multi_ractor_mode {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        // Assume single-ractor mode.
+                        self.push_insn(block, Insn::PatchPoint(Invariant::SingleRactorMode));
+                        // Invalidate output code on any constant writes associated with constants
+                        // referenced after the PatchPoint.
+                        self.push_insn(block, Insn::PatchPoint(Invariant::StableConstantNames { idlist }));
+                        let replacement = self.push_insn(block, Insn::Const { val: Const::Value(unsafe { (*ice).value }) });
+                        self.make_equal_to(insn_id, replacement);
+                    }
                     _ => { self.push_insn_id(block, insn_id); }
                 }
             }
@@ -1714,7 +1756,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     state.stack_push(fun.push_insn(block, Insn::Defined { op_type, obj, pushval, v }));
                 }
                 YARVINSN_opt_getconstant_path => {
-                    let ic = get_arg(pc, 0).as_ptr::<u8>();
+                    let ic = get_arg(pc, 0).as_ptr();
                     state.stack_push(fun.push_insn(block, Insn::GetConstantPath { ic }));
                 }
                 YARVINSN_branchunless => {
@@ -3565,4 +3607,70 @@ mod opt_tests {
               Return v7
         "#]]);
     }
+
+    #[test]
+    fn dont_replace_get_constant_path_with_empty_ic() {
+        eval("
+            def test = Kernel
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0():
+              v1:BasicObject = GetConstantPath 0x1000
+              Return v1
+        "#]]);
+    }
+
+    #[test]
+    fn dont_replace_get_constant_path_with_invalidated_ic() {
+        eval("
+            def test = Kernel
+            test
+            Kernel = 5
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0():
+              v1:BasicObject = GetConstantPath 0x1000
+              Return v1
+        "#]]);
+    }
+
+    #[test]
+    fn replace_get_constant_path_with_const() {
+        eval("
+            def test = Kernel
+            test
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0():
+              PatchPoint SingleRactorMode
+              PatchPoint StableConstantNames(0x1000, Kernel)
+              v5:BasicObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+              Return v5
+        "#]]);
+    }
+
+    #[test]
+    fn replace_nested_get_constant_path_with_const() {
+        eval("
+            module Foo
+              module Bar
+                class C
+                end
+              end
+            end
+            def test = Foo::Bar::C
+            test
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0():
+              PatchPoint SingleRactorMode
+              PatchPoint StableConstantNames(0x1000, Foo::Bar::C)
+              v5:BasicObject[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+              Return v5
+        "#]]);
+    }
 }