4 files changed, 126 insertions, 2 deletions

diff --git a/zjit.c b/zjit.c
index 5ddeb4aac6..32ccf521ad 100644
--- a/zjit.c
+++ b/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)
 {
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index 1ce9008af1..5372ead3e5 100644
--- a/zjit/bindgen/src/main.rs
+++ b/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")
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index ce1f6114cb..169789ef77 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/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(
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 17f7cb0e54..24b7cb6a0e 100644
--- a/zjit/src/hir.rs
+++ b/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
+ "#]]);
+ }
 }