ZJIT: Inline Thread.current (#15272)

Add `LoadEC` then it's just two `LoadField`.

5 files changed, 42 insertions, 3 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 06f991c738..b95d137222 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -457,6 +457,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::ArrayExtend { left, right, state } => { no_output!(gen_array_extend(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state))) },
         &Insn::GuardShape { val, shape, state } => gen_guard_shape(jit, asm, opnd!(val), shape, &function.frame_state(state)),
         Insn::LoadPC => gen_load_pc(asm),
+        Insn::LoadEC => gen_load_ec(),
         Insn::LoadSelf => gen_load_self(),
         &Insn::LoadField { recv, id, offset, return_type: _ } => gen_load_field(asm, opnd!(recv), id, offset),
         &Insn::StoreField { recv, id, offset, val } => no_output!(gen_store_field(asm, opnd!(recv), id, offset, opnd!(val))),
@@ -1041,6 +1042,10 @@ fn gen_load_pc(asm: &mut Assembler) -> Opnd {
     asm.load(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_PC))
 }
 
+fn gen_load_ec() -> Opnd {
+    EC
+}
+
 fn gen_load_self() -> Opnd {
     Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF)
 }
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
index 443ed0d86e..83919e5369 100644
--- a/zjit/src/cruby.rs
+++ b/zjit/src/cruby.rs
@@ -1377,6 +1377,8 @@ pub(crate) mod ids {
         name: aref               content: b"[]"
         name: len
         name: _as_heap
+        name: thread_ptr
+        name: self_              content: b"self"
     }
 
     /// Get an CRuby `ID` to an interned string, e.g. a particular method name.
diff --git a/zjit/src/cruby_methods.rs b/zjit/src/cruby_methods.rs
index 190ac52eac..d1f76e8da0 100644
--- a/zjit/src/cruby_methods.rs
+++ b/zjit/src/cruby_methods.rs
@@ -158,6 +158,7 @@ pub fn init() -> Annotations {
         ($module:ident, $method_name:literal, $inline:ident) => {
             let mut props = FnProperties::default();
             props.inline = $inline;
+            #[allow(unused_unsafe)]
             annotate_c_method(cfuncs, unsafe { $module }, $method_name, props);
         };
         ($module:ident, $method_name:literal, $inline:ident, $return_type:expr $(, $properties:ident)*) => {
@@ -167,6 +168,7 @@ pub fn init() -> Annotations {
             $(
                 props.$properties = true;
             )*
+            #[allow(unused_unsafe)]
             annotate_c_method(cfuncs, unsafe { $module }, $method_name, props);
         };
         ($module:ident, $method_name:literal, $return_type:expr $(, $properties:ident)*) => {
@@ -240,7 +242,7 @@ pub fn init() -> Annotations {
     annotate!(rb_cInteger, "<=", inline_integer_le);
     annotate!(rb_cString, "to_s", inline_string_to_s, types::StringExact);
     let thread_singleton = unsafe { rb_singleton_class(rb_cThread) };
-    annotate!(thread_singleton, "current", types::BasicObject, no_gc, leaf);
+    annotate!(thread_singleton, "current", inline_thread_current, types::BasicObject, no_gc, leaf);
 
     annotate_builtin!(rb_mKernel, "Float", types::Float);
     annotate_builtin!(rb_mKernel, "Integer", types::Integer);
@@ -269,6 +271,26 @@ fn inline_string_to_s(fun: &mut hir::Function, block: hir::BlockId, recv: hir::I
     None
 }
 
+fn inline_thread_current(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], _state: hir::InsnId) -> Option<hir::InsnId> {
+    let &[] = args else { return None; };
+    let ec = fun.push_insn(block, hir::Insn::LoadEC);
+    let thread_ptr = fun.push_insn(block, hir::Insn::LoadField {
+        recv: ec,
+        id: ID!(thread_ptr),
+        offset: RUBY_OFFSET_EC_THREAD_PTR as i32,
+        return_type: types::CPtr,
+    });
+    let thread_self = fun.push_insn(block, hir::Insn::LoadField {
+        recv: thread_ptr,
+        id: ID!(self_),
+        offset: RUBY_OFFSET_THREAD_SELF as i32,
+        // TODO(max): Add Thread type. But Thread.current is not guaranteed to be an exact Thread.
+        // You can make subclasses...
+        return_type: types::BasicObject,
+    });
+    Some(thread_self)
+}
+
 fn inline_kernel_itself(_fun: &mut hir::Function, _block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], _state: hir::InsnId) -> Option<hir::InsnId> {
     if args.is_empty() {
         // No need to coerce the receiver; that is done by the SendWithoutBlock rewriting.
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 40c6092e56..333d5e5bff 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -736,6 +736,8 @@ pub enum Insn {
 
     /// Load cfp->pc
     LoadPC,
+    /// Load EC
+    LoadEC,
     /// Load cfp->self
     LoadSelf,
     LoadField { recv: InsnId, id: ID, offset: i32, return_type: Type },
@@ -980,6 +982,7 @@ impl Insn {
             Insn::GetLocal   { .. } => false,
             Insn::IsNil      { .. } => false,
             Insn::LoadPC => false,
+            Insn::LoadEC => false,
             Insn::LoadSelf => false,
             Insn::LoadField { .. } => false,
             Insn::CCall { elidable, .. } => !elidable,
@@ -1272,6 +1275,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::DefinedIvar { self_val, id, .. } => write!(f, "DefinedIvar {self_val}, :{}", id.contents_lossy()),
             Insn::GetIvar { self_val, id, .. } => write!(f, "GetIvar {self_val}, :{}", id.contents_lossy()),
             Insn::LoadPC => write!(f, "LoadPC"),
+            Insn::LoadEC => write!(f, "LoadEC"),
             Insn::LoadSelf => write!(f, "LoadSelf"),
             &Insn::LoadField { recv, id, offset, return_type: _ } => write!(f, "LoadField {recv}, :{}@{:p}", id.contents_lossy(), self.ptr_map.map_offset(offset)),
             &Insn::StoreField { recv, id, offset, val } => write!(f, "StoreField {recv}, :{}@{:p}, {val}", id.contents_lossy(), self.ptr_map.map_offset(offset)),
@@ -1775,6 +1779,7 @@ impl Function {
                     | SideExit {..}
                     | EntryPoint {..}
                     | LoadPC
+                    | LoadEC
                     | LoadSelf
                     | IncrCounterPtr {..}
                     | IncrCounter(_)) => result.clone(),
@@ -2069,6 +2074,7 @@ impl Function {
             Insn::GetGlobal { .. } => types::BasicObject,
             Insn::GetIvar { .. } => types::BasicObject,
             Insn::LoadPC => types::CPtr,
+            Insn::LoadEC => types::CPtr,
             Insn::LoadSelf => types::BasicObject,
             &Insn::LoadField { return_type, .. } => return_type,
             Insn::GetSpecialSymbol { .. } => types::BasicObject,
@@ -3397,6 +3403,7 @@ impl Function {
             | &Insn::Param
             | &Insn::EntryPoint { .. }
             | &Insn::LoadPC
+            | &Insn::LoadEC
             | &Insn::LoadSelf
             | &Insn::GetLocal { .. }
             | &Insn::PutSpecialObject { .. }
@@ -4101,6 +4108,7 @@ impl Function {
             | Insn::IsBlockGiven
             | Insn::GetGlobal { .. }
             | Insn::LoadPC
+            | Insn::LoadEC
             | Insn::LoadSelf
             | Insn::Snapshot { .. }
             | Insn::Jump { .. }
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index b1ab8a0605..82f54f611a 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -5965,10 +5965,12 @@ mod hir_opt_tests {
           v20:Class[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           PatchPoint MethodRedefined(Class@0x1010, current@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(Class@0x1010)
+          v24:CPtr = LoadEC
+          v25:CPtr = LoadField v24, :thread_ptr@0x1048
+          v26:BasicObject = LoadField v25, :self@0x1049
           IncrCounter inline_cfunc_optimized_send_count
-          v25:BasicObject = CCall Thread.current@0x1048, v20
           CheckInterrupts
-          Return v25
+          Return v26
         ");
     }