YJIT: Avoid creating payloads for non-JITed ISEQs (#6549)

* YJIT: Count freed ISEQs

* YJIT: Avoid creating payloads for non-JITed ISEQs

5 files changed, 52 insertions, 28 deletions

diff --git a/yjit.rb b/yjit.rb
index 226f2a8134..b80861dbfb 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -214,6 +214,7 @@ module RubyVM::YJIT
       $stderr.puts "compilation_failure:   " + ("%10d" % compilation_failure) if compilation_failure != 0
       $stderr.puts "compiled_iseq_count:   " + ("%10d" % stats[:compiled_iseq_count])
       $stderr.puts "compiled_block_count:  " + ("%10d" % stats[:compiled_block_count])
+      $stderr.puts "freed_iseq_count:      " + ("%10d" % stats[:freed_iseq_count])
       $stderr.puts "invalidation_count:    " + ("%10d" % stats[:invalidation_count])
       $stderr.puts "constant_state_bumps:  " + ("%10d" % stats[:constant_state_bumps])
       $stderr.puts "inline_code_size:      " + ("%10d" % stats[:inline_code_size])
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index c8078bb6e3..1805a2bb13 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -491,14 +491,14 @@ impl IseqPayload {
 
 /// Get the payload for an iseq. For safety it's up to the caller to ensure the returned `&mut`
 /// upholds aliasing rules and that the argument is a valid iseq.
-pub unsafe fn load_iseq_payload(iseq: IseqPtr) -> Option<&'static mut IseqPayload> {
-    let payload = rb_iseq_get_yjit_payload(iseq);
+pub fn get_iseq_payload(iseq: IseqPtr) -> Option<&'static mut IseqPayload> {
+    let payload = unsafe { rb_iseq_get_yjit_payload(iseq) };
     let payload: *mut IseqPayload = payload.cast();
-    payload.as_mut()
+    unsafe { payload.as_mut() }
 }
 
 /// Get the payload object associated with an iseq. Create one if none exists.
-fn get_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
+fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
     type VoidPtr = *mut c_void;
 
     let payload_non_null = unsafe {
@@ -546,6 +546,9 @@ pub extern "C" fn rb_yjit_iseq_free(payload: *mut c_void) {
     // SAFETY: We got the pointer from Box::into_raw().
     let payload = unsafe { Box::from_raw(payload) };
 
+    // Increment the freed iseq count
+    incr_counter!(freed_iseq_count);
+
     // Remove all blocks in the payload from global invariants table.
     for versions in &payload.version_map {
         for block in versions {
@@ -679,8 +682,19 @@ pub extern "C" fn rb_yjit_iseq_update_references(payload: *mut c_void) {
 }
 
 /// Get all blocks for a particular place in an iseq.
-fn get_version_list(blockid: BlockId) -> &'static mut VersionList {
-    let payload = get_iseq_payload(blockid.iseq);
+fn get_version_list(blockid: BlockId) -> Option<&'static mut VersionList> {
+    let insn_idx = blockid.idx.as_usize();
+    match get_iseq_payload(blockid.iseq) {
+        Some(payload) if insn_idx < payload.version_map.len() => {
+            Some(payload.version_map.get_mut(insn_idx).unwrap())
+        },
+        _ => None
+    }
+}
+
+/// Get or create all blocks for a particular place in an iseq.
+fn get_or_create_version_list(blockid: BlockId) -> &'static mut VersionList {
+    let payload = get_or_create_iseq_payload(blockid.iseq);
     let insn_idx = blockid.idx.as_usize();
 
     // Expand the version map as necessary
@@ -695,32 +709,34 @@ fn get_version_list(blockid: BlockId) -> &'static mut VersionList {
 
 /// Take all of the blocks for a particular place in an iseq
 pub fn take_version_list(blockid: BlockId) -> VersionList {
-    let payload = get_iseq_payload(blockid.iseq);
     let insn_idx = blockid.idx.as_usize();
-
-    if insn_idx >= payload.version_map.len() {
-        VersionList::default()
-    } else {
-        mem::take(&mut payload.version_map[insn_idx])
+    match get_iseq_payload(blockid.iseq) {
+        Some(payload) if insn_idx < payload.version_map.len() => {
+            mem::take(&mut payload.version_map[insn_idx])
+        },
+        _ => VersionList::default(),
     }
 }
 
 /// Count the number of block versions matching a given blockid
 fn get_num_versions(blockid: BlockId) -> usize {
     let insn_idx = blockid.idx.as_usize();
-    let payload = get_iseq_payload(blockid.iseq);
-
-    payload
-        .version_map
-        .get(insn_idx)
-        .map(|versions| versions.len())
-        .unwrap_or(0)
+    match get_iseq_payload(blockid.iseq) {
+        Some(payload) => {
+            payload
+                .version_map
+                .get(insn_idx)
+                .map(|versions| versions.len())
+                .unwrap_or(0)
+        }
+        None => 0,
+    }
 }
 
-/// Get a list of block versions generated for an iseq
+/// Get or create a list of block versions generated for an iseq
 /// This is used for disassembly (see disasm.rs)
-pub fn get_iseq_block_list(iseq: IseqPtr) -> Vec<BlockRef> {
-    let payload = get_iseq_payload(iseq);
+pub fn get_or_create_iseq_block_list(iseq: IseqPtr) -> Vec<BlockRef> {
+    let payload = get_or_create_iseq_payload(iseq);
 
     let mut blocks = Vec::<BlockRef>::new();
 
@@ -741,7 +757,10 @@ pub fn get_iseq_block_list(iseq: IseqPtr) -> Vec<BlockRef> {
 /// Retrieve a basic block version for an (iseq, idx) tuple
 /// This will return None if no version is found
 fn find_block_version(blockid: BlockId, ctx: &Context) -> Option<BlockRef> {
-    let versions = get_version_list(blockid);
+    let versions = match get_version_list(blockid) {
+        Some(versions) => versions,
+        None => return None,
+    };
 
     // Best match found
     let mut best_version: Option<BlockRef> = None;
@@ -802,7 +821,7 @@ fn add_block_version(blockref: &BlockRef, cb: &CodeBlock) {
     // Function entry blocks must have stack size 0
     assert!(!(block.blockid.idx == 0 && block.ctx.stack_size > 0));
 
-    let version_list = get_version_list(block.blockid);
+    let version_list = get_or_create_version_list(block.blockid);
 
     version_list.push(blockref.clone());
 
@@ -830,7 +849,10 @@ fn add_block_version(blockref: &BlockRef, cb: &CodeBlock) {
 /// Remove a block version from the version map of its parent ISEQ
 fn remove_block_version(blockref: &BlockRef) {
     let block = blockref.borrow();
-    let version_list = get_version_list(block.blockid);
+    let version_list = match get_version_list(block.blockid) {
+        Some(version_list) => version_list,
+        None => return,
+    };
 
     // Retain the versions that are not this one
     version_list.retain(|other| blockref != other);
diff --git a/yjit/src/disasm.rs b/yjit/src/disasm.rs
index 10a89bafd0..9e45dffd60 100644
--- a/yjit/src/disasm.rs
+++ b/yjit/src/disasm.rs
@@ -42,7 +42,7 @@ pub fn disasm_iseq_insn_range(iseq: IseqPtr, start_idx: u32, end_idx: u32) -> St
     let mut out = String::from("");
 
     // Get a list of block versions generated for this iseq
-    let mut block_list = get_iseq_block_list(iseq);
+    let mut block_list = get_or_create_iseq_block_list(iseq);
 
     // Get a list of codeblocks relevant to this iseq
     let global_cb = crate::codegen::CodegenGlobals::get_inline_cb();
@@ -206,7 +206,7 @@ fn insns_compiled(iseq: IseqPtr) -> Vec<(String, u32)> {
     let mut insn_vec = Vec::new();
 
     // Get a list of block versions generated for this iseq
-    let block_list = get_iseq_block_list(iseq);
+    let block_list = get_or_create_iseq_block_list(iseq);
 
     // For each block associated with this iseq
     for blockref in &block_list {
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index c7c0701e74..07de3374c8 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -535,7 +535,7 @@ pub extern "C" fn rb_yjit_tracing_invalidate_all() {
         unsafe { rb_yjit_for_each_iseq(Some(invalidate_all_blocks_for_tracing)) };
 
         extern "C" fn invalidate_all_blocks_for_tracing(iseq: IseqPtr) {
-            if let Some(payload) = unsafe { load_iseq_payload(iseq) } {
+            if let Some(payload) = unsafe { get_iseq_payload(iseq) } {
                 // C comment:
                 //   Leaking the blocks for now since we might have situations where
                 //   a different ractor is waiting for the VM lock in branch_stub_hit().
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index bf968ee563..0ad77fc5df 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -252,6 +252,7 @@ make_counters! {
     compiled_iseq_count,
     compiled_block_count,
     compilation_failure,
+    freed_iseq_count,
 
     exit_from_branch_stub,