ZJIT: Carve out IseqPayload into a separate module (ruby#15098)

Author: k0kubun

zjit/src/codegen.rs

@@ -13,7 +13,8 @@ use crate::invariants::{
     track_bop_assumption, track_cme_assumption, track_no_ep_escape_assumption, track_no_trace_point_assumption,
     track_single_ractor_assumption, track_stable_constant_names_assumption, track_no_singleton_class_assumption
 };
-use crate::gc::{append_gc_offsets, get_or_create_iseq_payload, get_or_create_iseq_payload_ptr, IseqCodePtrs, IseqPayload, IseqStatus};
+use crate::gc::append_gc_offsets;
+use crate::payload::{get_or_create_iseq_payload, get_or_create_iseq_payload_ptr, IseqCodePtrs, IseqPayload, IseqStatus};
 use crate::state::ZJITState;
 use crate::stats::{send_fallback_counter, exit_counter_for_compile_error, incr_counter, incr_counter_by, send_fallback_counter_for_method_type, send_without_block_fallback_counter_for_method_type, send_without_block_fallback_counter_for_optimized_method_type, send_fallback_counter_ptr_for_opcode, CompileError};
 use crate::stats::{counter_ptr, with_time_stat, Counter, Counter::{compile_time_ns, exit_compile_error}};

zjit/src/gc.rs

@@ -1,94 +1,11 @@
 //! This module is responsible for marking/moving objects on GC.
 
 use std::{ffi::c_void, ops::Range};
-use crate::codegen::IseqCallRef;
-use crate::stats::CompileError;
-use crate::{cruby::*, profile::IseqProfile, state::ZJITState, stats::with_time_stat, virtualmem::CodePtr};
+use crate::{cruby::*, state::ZJITState, stats::with_time_stat, virtualmem::CodePtr};
+use crate::payload::{IseqPayload, get_or_create_iseq_payload, payload_ptr_as_mut};
 use crate::stats::Counter::gc_time_ns;
 use crate::state::gc_mark_raw_samples;
 
-/// This is all the data ZJIT stores on an ISEQ. We mark objects in this struct on GC.
-#[derive(Debug)]
-pub struct IseqPayload {
-    /// Compilation status of the ISEQ. It has the JIT code address of the first block if Compiled.
-    pub status: IseqStatus,
-
-    /// Type information of YARV instruction operands
-    pub profile: IseqProfile,
-
-    /// GC offsets of the JIT code. These are the addresses of objects that need to be marked.
-    pub gc_offsets: Vec<CodePtr>,
-
-    /// JIT-to-JIT calls in the ISEQ. The IseqPayload's ISEQ is the caller of it.
-    pub iseq_calls: Vec<IseqCallRef>,
-}
-
-impl IseqPayload {
-    fn new(iseq_size: u32) -> Self {
-        Self {
-            status: IseqStatus::NotCompiled,
-            profile: IseqProfile::new(iseq_size),
-            gc_offsets: vec![],
-            iseq_calls: vec![],
-        }
-    }
-}
-
-/// Set of CodePtrs for an ISEQ
-#[derive(Clone, Debug, PartialEq)]
-pub struct IseqCodePtrs {
-    /// Entry for the interpreter
-    pub start_ptr: CodePtr,
-    /// Entries for JIT-to-JIT calls
-    pub jit_entry_ptrs: Vec<CodePtr>,
-}
-
-#[derive(Debug, PartialEq)]
-pub enum IseqStatus {
-    Compiled(IseqCodePtrs),
-    CantCompile(CompileError),
-    NotCompiled,
-}
-
-/// Get a pointer to the payload object associated with an ISEQ. Create one if none exists.
-pub fn get_or_create_iseq_payload_ptr(iseq: IseqPtr) -> *mut IseqPayload {
-    type VoidPtr = *mut c_void;
-
-    unsafe {
-        let payload = rb_iseq_get_zjit_payload(iseq);
-        if payload.is_null() {
-            // Allocate a new payload with Box and transfer ownership to the GC.
-            // We drop the payload with Box::from_raw when the GC frees the ISEQ and calls us.
-            // NOTE(alan): Sometimes we read from an ISEQ without ever writing to it.
-            // We allocate in those cases anyways.
-            let iseq_size = get_iseq_encoded_size(iseq);
-            let new_payload = IseqPayload::new(iseq_size);
-            let new_payload = Box::into_raw(Box::new(new_payload));
-            rb_iseq_set_zjit_payload(iseq, new_payload as VoidPtr);
-
-            new_payload
-        } else {
-            payload as *mut IseqPayload
-        }
-    }
-}
-
-/// Get the payload object associated with an ISEQ. Create one if none exists.
-pub fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
-    let payload_non_null = get_or_create_iseq_payload_ptr(iseq);
-    payload_ptr_as_mut(payload_non_null)
-}
-
-/// Convert an IseqPayload pointer to a mutable reference. Only one reference
-/// should be kept at a time.
-fn payload_ptr_as_mut(payload_ptr: *mut IseqPayload) -> &'static mut IseqPayload {
-    // SAFETY: we should have the VM lock and all other Ruby threads should be asleep. So we have
-    // exclusive mutable access.
-    // Hmm, nothing seems to stop calling this on the same
-    // iseq twice, though, which violates aliasing rules.
-    unsafe { payload_ptr.as_mut() }.unwrap()
-}
-
 /// GC callback for marking GC objects in the per-ISEQ payload.
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_iseq_mark(payload: *mut c_void) {

zjit/src/hir.rs

@@ -6,7 +6,7 @@
 #![allow(clippy::if_same_then_else)]
 #![allow(clippy::match_like_matches_macro)]
 use crate::{
-    cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, gc::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState
+    cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState
 };
 use std::{
     cell::RefCell, collections::{HashMap, HashSet, VecDeque}, ffi::{c_void, c_uint, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter

zjit/src/invariants.rs

@@ -2,7 +2,8 @@
 
 use std::{collections::{HashMap, HashSet}, mem};
 
-use crate::{backend::lir::{asm_comment, Assembler}, cruby::{iseq_name, rb_callable_method_entry_t, rb_gc_location, ruby_basic_operators, src_loc, with_vm_lock, IseqPtr, RedefinitionFlag, ID, VALUE}, gc::IseqPayload, hir::Invariant, options::debug, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
+use crate::{backend::lir::{asm_comment, Assembler}, cruby::{iseq_name, rb_callable_method_entry_t, rb_gc_location, ruby_basic_operators, src_loc, with_vm_lock, IseqPtr, RedefinitionFlag, ID, VALUE}, hir::Invariant, options::debug, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
+use crate::payload::IseqPayload;
 use crate::stats::with_time_stat;
 use crate::stats::Counter::invalidation_time_ns;
 use crate::gc::remove_gc_offsets;
@@ -367,7 +368,7 @@ pub fn track_no_trace_point_assumption(patch_point_ptr: CodePtr, side_exit_ptr:
 
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_tracing_invalidate_all() {
-    use crate::gc::{get_or_create_iseq_payload, IseqStatus};
+    use crate::payload::{get_or_create_iseq_payload, IseqStatus};
     use crate::cruby::{for_each_iseq, rb_iseq_reset_jit_func};
 
     if !zjit_enabled_p() {

zjit/src/lib.rs

@@ -28,3 +28,4 @@ mod profile;
 mod invariants;
 mod bitset;
 mod gc;
+mod payload;

zjit/src/payload.rs

@@ -0,0 +1,86 @@
+use std::ffi::c_void;
+use crate::codegen::IseqCallRef;
+use crate::stats::CompileError;
+use crate::{cruby::*, profile::IseqProfile, virtualmem::CodePtr};
+
+/// This is all the data ZJIT stores on an ISEQ. We mark objects in this struct on GC.
+#[derive(Debug)]
+pub struct IseqPayload {
+    /// Compilation status of the ISEQ. It has the JIT code address of the first block if Compiled.
+    pub status: IseqStatus,
+
+    /// Type information of YARV instruction operands
+    pub profile: IseqProfile,
+
+    /// GC offsets of the JIT code. These are the addresses of objects that need to be marked.
+    pub gc_offsets: Vec<CodePtr>,
+
+    /// JIT-to-JIT calls in the ISEQ. The IseqPayload's ISEQ is the caller of it.
+    pub iseq_calls: Vec<IseqCallRef>,
+}
+
+impl IseqPayload {
+    fn new(iseq_size: u32) -> Self {
+        Self {
+            status: IseqStatus::NotCompiled,
+            profile: IseqProfile::new(iseq_size),
+            gc_offsets: vec![],
+            iseq_calls: vec![],
+        }
+    }
+}
+
+/// Set of CodePtrs for an ISEQ
+#[derive(Clone, Debug, PartialEq)]
+pub struct IseqCodePtrs {
+    /// Entry for the interpreter
+    pub start_ptr: CodePtr,
+    /// Entries for JIT-to-JIT calls
+    pub jit_entry_ptrs: Vec<CodePtr>,
+}
+
+#[derive(Debug, PartialEq)]
+pub enum IseqStatus {
+    Compiled(IseqCodePtrs),
+    CantCompile(CompileError),
+    NotCompiled,
+}
+
+/// Get a pointer to the payload object associated with an ISEQ. Create one if none exists.
+pub fn get_or_create_iseq_payload_ptr(iseq: IseqPtr) -> *mut IseqPayload {
+    type VoidPtr = *mut c_void;
+
+    unsafe {
+        let payload = rb_iseq_get_zjit_payload(iseq);
+        if payload.is_null() {
+            // Allocate a new payload with Box and transfer ownership to the GC.
+            // We drop the payload with Box::from_raw when the GC frees the ISEQ and calls us.
+            // NOTE(alan): Sometimes we read from an ISEQ without ever writing to it.
+            // We allocate in those cases anyways.
+            let iseq_size = get_iseq_encoded_size(iseq);
+            let new_payload = IseqPayload::new(iseq_size);
+            let new_payload = Box::into_raw(Box::new(new_payload));
+            rb_iseq_set_zjit_payload(iseq, new_payload as VoidPtr);
+
+            new_payload
+        } else {
+            payload as *mut IseqPayload
+        }
+    }
+}
+
+/// Get the payload object associated with an ISEQ. Create one if none exists.
+pub fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
+    let payload_non_null = get_or_create_iseq_payload_ptr(iseq);
+    payload_ptr_as_mut(payload_non_null)
+}
+
+/// Convert an IseqPayload pointer to a mutable reference. Only one reference
+/// should be kept at a time.
+pub fn payload_ptr_as_mut(payload_ptr: *mut IseqPayload) -> &'static mut IseqPayload {
+    // SAFETY: we should have the VM lock and all other Ruby threads should be asleep. So we have
+    // exclusive mutable access.
+    // Hmm, nothing seems to stop calling this on the same
+    // iseq twice, though, which violates aliasing rules.
+    unsafe { payload_ptr.as_mut() }.unwrap()
+}

zjit/src/profile.rs

@@ -3,7 +3,7 @@
 // We use the YARV bytecode constants which have a CRuby-style name
 #![allow(non_upper_case_globals)]
 
-use crate::{cruby::*, gc::get_or_create_iseq_payload, options::{get_option, NumProfiles}};
+use crate::{cruby::*, payload::get_or_create_iseq_payload, options::{get_option, NumProfiles}};
 use crate::distribution::{Distribution, DistributionSummary};
 use crate::stats::Counter::profile_time_ns;
 use crate::stats::with_time_stat;