YJIT: reject large stacks so we can use i8/u8 stack_size and stack_offset (#7412)

* Reject large stacks so we can use i8/u8 stack_size and stack_offset

* Add rejection test for iseq too long as well

3 files changed, 23 insertions, 0 deletions

diff --git a/yjit.rb b/yjit.rb
index 4ba16c6b26..459d58f918 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -261,6 +261,8 @@ module RubyVM::YJIT
         $stderr.puts "num_send_x86_reg:      " + format_number(13, stats[:num_send_x86_reg])
       end
 
+      $stderr.puts "iseq_stack_too_large:  " + format_number(13, stats[:iseq_stack_too_large])
+      $stderr.puts "iseq_too_long:         " + format_number(13, stats[:iseq_too_long])
       $stderr.puts "bindings_allocations:  " + format_number(13, stats[:binding_allocations])
       $stderr.puts "bindings_set:          " + format_number(13, stats[:binding_set])
       $stderr.puts "compilation_failure:   " + format_number(13, compilation_failure) if compilation_failure != 0
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 5d32f613c8..2587042d79 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -348,6 +348,9 @@ make_counters! {
     num_send_polymorphic,
     num_send_x86_rel32,
     num_send_x86_reg,
+
+    iseq_stack_too_large,
+    iseq_too_long,
 }
 
 //===========================================================================
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index bfd3060f46..e133049ee8 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -4,6 +4,7 @@ use crate::cruby::*;
 use crate::invariants::*;
 use crate::options::*;
 use crate::stats::YjitExitLocations;
+use crate::stats::incr_counter;
 
 use std::os::raw;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -96,6 +97,23 @@ fn rb_bug_panic_hook() {
 /// NOTE: this should be wrapped in RB_VM_LOCK_ENTER(), rb_vm_barrier() on the C side
 #[no_mangle]
 pub extern "C" fn rb_yjit_iseq_gen_entry_point(iseq: IseqPtr, ec: EcPtr) -> *const u8 {
+    // Reject ISEQs with very large temp stacks,
+    // this will allow us to use u8/i8 values to track stack_size and sp_offset
+    let stack_max = unsafe { rb_get_iseq_body_stack_max(iseq) };
+    if stack_max >= i8::MAX as u32 {
+        incr_counter!(iseq_stack_too_large);
+        return std::ptr::null();
+    }
+
+    // Reject ISEQs that are too long,
+    // this will allow us to use u16 for instruction indices if we want to,
+    // very long ISEQs are also much more likely to be initialization code
+    let iseq_size = unsafe { get_iseq_encoded_size(iseq) };
+    if iseq_size >= u16::MAX as u32 {
+        incr_counter!(iseq_too_long);
+        return std::ptr::null();
+    }
+
     let maybe_code_ptr = gen_entry_point(iseq, ec);
 
     match maybe_code_ptr {