ZJIT: Allow testing JIT code on zjit-test (ruby#14639)

* ZJIT: Allow testing JIT code on zjit-test

* Resurrect TestingAllocator tests

Author: k0kubun

zjit/src/asm/mod.rs

@@ -317,19 +317,13 @@ impl fmt::LowerHex for CodeBlock {
 impl CodeBlock {
     /// Stubbed CodeBlock for testing. Can't execute generated code.
     pub fn new_dummy() -> Self {
-        const DEFAULT_MEM_SIZE: usize = 1024;
+        const DEFAULT_MEM_SIZE: usize = 1024 * 1024;
         CodeBlock::new_dummy_sized(DEFAULT_MEM_SIZE)
     }
 
     pub fn new_dummy_sized(mem_size: usize) -> Self {
-        use std::ptr::NonNull;
         use crate::virtualmem::*;
-        use crate::virtualmem::tests::TestingAllocator;
-
-        let alloc = TestingAllocator::new(mem_size);
-        let mem_start: *const u8 = alloc.mem_start();
-        let virt_mem = VirtualMem::new(alloc, 1, NonNull::new(mem_start as *mut u8).unwrap(), mem_size, 128 * 1024 * 1024);
-
+        let virt_mem = VirtualMem::alloc(mem_size, None);
         Self::new(Rc::new(RefCell::new(virt_mem)), false)
     }
 }

zjit/src/backend/arm64/mod.rs

@@ -2058,10 +2058,12 @@ mod tests {
 
         // `IMMEDIATE_MAX_VALUE` number of dummy instructions will be generated
         // plus a compare, a jump instruction, and a label.
-        const MEMORY_REQUIRED: usize = (IMMEDIATE_MAX_VALUE + 8)*4;
+        // Adding page_size to avoid OOM on the last page.
+        let page_size = unsafe { rb_jit_get_page_size() } as usize;
+        let memory_required = (IMMEDIATE_MAX_VALUE + 8) * 4 + page_size;
 
         let mut asm = Assembler::new();
-        let mut cb = CodeBlock::new_dummy_sized(MEMORY_REQUIRED);
+        let mut cb = CodeBlock::new_dummy_sized(memory_required);
 
         let far_label = asm.new_label("far");
 

zjit/src/codegen.rs

@@ -75,11 +75,6 @@ impl JITState {
 /// See jit_compile_exception() for details.
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, jit_exception: bool) -> *const u8 {
-    // Do not test the JIT code in HIR tests
-    if cfg!(test) {
-        return std::ptr::null();
-    }
-
     // Take a lock to avoid writing to ISEQ in parallel with Ractors.
     // with_vm_lock() does nothing if the program doesn't use Ractors.
     with_vm_lock(src_loc!(), || {

zjit/src/cruby.rs

@@ -999,7 +999,7 @@ pub use manual_defs::*;
 pub mod test_utils {
     use std::{ptr::null, sync::Once};
 
-    use crate::{options::{internal_set_num_profiles, rb_zjit_call_threshold, rb_zjit_prepare_options, DEFAULT_CALL_THRESHOLD}, state::{rb_zjit_enabled_p, ZJITState}};
+    use crate::{options::{rb_zjit_call_threshold, rb_zjit_prepare_options, set_call_threshold, DEFAULT_CALL_THRESHOLD}, state::{rb_zjit_enabled_p, ZJITState}};
 
     use super::*;
 
@@ -1026,7 +1026,7 @@ pub mod test_utils {
 
             // The default rb_zjit_profile_threshold is too high, so lower it for HIR tests.
             if rb_zjit_call_threshold == DEFAULT_CALL_THRESHOLD {
-                internal_set_num_profiles(1);
+                set_call_threshold(2);
             }
 
             // Pass command line options so the VM loads core library methods defined in
@@ -1099,6 +1099,12 @@ pub mod test_utils {
         })
     }
 
+    /// Get the #inspect of a given Ruby program in Rust string
+    pub fn inspect(program: &str) -> String {
+        let inspect = format!("({program}).inspect");
+        ruby_str_to_rust_string(eval(&inspect))
+    }
+
     /// Get the ISeq of a specified method
     pub fn get_method_iseq(recv: &str, name: &str) -> *const rb_iseq_t {
         let wrapped_iseq = eval(&format!("RubyVM::InstructionSequence.of({}.method(:{}))", recv, name));

zjit/src/hir.rs

@@ -6545,6 +6545,7 @@ mod graphviz_tests {
 #[cfg(test)]
 mod opt_tests {
     use super::*;
+    use crate::options::*;
     use insta::assert_snapshot;
 
     #[track_caller]
@@ -8481,6 +8482,7 @@ mod opt_tests {
           CheckInterrupts
           Return v37
         ");
+        assert_snapshot!(inspect("test"), @"{}");
     }
 
     #[test]
@@ -9580,8 +9582,7 @@ mod opt_tests {
 
     #[test]
     fn test_dont_optimize_getivar_polymorphic() {
-        crate::options::rb_zjit_prepare_options();
-        crate::options::internal_set_num_profiles(3);
+        set_call_threshold(3);
         eval("
             class C
               attr_reader :foo, :bar

zjit/src/options.rs

@@ -301,12 +301,11 @@ fn update_profile_threshold() {
     }
 }
 
+/// Update --zjit-call-threshold for testing
 #[cfg(test)]
-pub fn internal_set_num_profiles(n: u8) {
-    let options = unsafe { OPTIONS.as_mut().unwrap() };
-    options.num_profiles = n;
-    let call_threshold = n.saturating_add(1);
+pub fn set_call_threshold(call_threshold: u64) {
     unsafe { rb_zjit_call_threshold = call_threshold as u64; }
+    rb_zjit_prepare_options();
     update_profile_threshold();
 }
 

zjit/src/state.rs

@@ -54,45 +54,17 @@ static mut ZJIT_STATE: Option<ZJITState> = None;
 impl ZJITState {
     /// Initialize the ZJIT globals
     pub fn init() {
-        #[cfg(not(test))]
         let mut cb = {
-            use crate::cruby::*;
             use crate::options::*;
-
-            let exec_mem_bytes: usize = get_option!(exec_mem_bytes);
-            let virt_block: *mut u8 = unsafe { rb_jit_reserve_addr_space(64 * 1024 * 1024) };
-
-            // Memory protection syscalls need page-aligned addresses, so check it here. Assuming
-            // `virt_block` is page-aligned, `second_half` should be page-aligned as long as the
-            // page size in bytes is a power of two 2¹⁹ or smaller. This is because the user
-            // requested size is half of mem_option × 2²⁰ as it's in MiB.
-            //
-            // Basically, we don't support x86-64 2MiB and 1GiB pages. ARMv8 can do up to 64KiB
-            // (2¹⁶ bytes) pages, which should be fine. 4KiB pages seem to be the most popular though.
-            let page_size = unsafe { rb_jit_get_page_size() };
-            assert_eq!(
-                virt_block as usize % page_size as usize, 0,
-                "Start of virtual address block should be page-aligned",
-            );
-
             use crate::virtualmem::*;
-            use std::ptr::NonNull;
             use std::rc::Rc;
             use std::cell::RefCell;
 
-            let mem_block = VirtualMem::new(
-                crate::virtualmem::sys::SystemAllocator {},
-                page_size,
-                NonNull::new(virt_block).unwrap(),
-                exec_mem_bytes,
-                get_option!(mem_bytes)
-            );
+            let mem_block = VirtualMem::alloc(get_option!(exec_mem_bytes), Some(get_option!(mem_bytes)));
             let mem_block = Rc::new(RefCell::new(mem_block));
 
             CodeBlock::new(mem_block.clone(), get_option!(dump_disasm))
         };
-        #[cfg(test)]
-        let mut cb = CodeBlock::new_dummy();
 
         let exit_trampoline = gen_exit_trampoline(&mut cb).unwrap();
         let function_stub_hit_trampoline = gen_function_stub_hit_trampoline(&mut cb).unwrap();

zjit/src/virtualmem.rs

@@ -4,18 +4,11 @@
 // benefit.
 
 use std::ptr::NonNull;
-
+use crate::cruby::*;
 use crate::stats::zjit_alloc_size;
 
-#[cfg(test)]
-use crate::options::get_option;
-
-#[cfg(not(test))]
 pub type VirtualMem = VirtualMemory<sys::SystemAllocator>;
 
-#[cfg(test)]
-pub type VirtualMem = VirtualMemory<tests::TestingAllocator>;
-
 /// Memory for generated executable machine code. When not testing, we reserve address space for
 /// the entire region upfront and map physical memory into the reserved address space as needed. On
 /// Linux, this is basically done using an `mmap` with `PROT_NONE` upfront and gradually using
@@ -33,7 +26,7 @@ pub struct VirtualMemory<A: Allocator> {
     region_size_bytes: usize,
 
     /// mapped_region_bytes + zjit_alloc_size may not increase beyond this limit.
-    memory_limit_bytes: usize,
+    memory_limit_bytes: Option<usize>,
 
     /// Number of bytes per "page", memory protection permission can only be controlled at this
     /// granularity.
@@ -113,14 +106,36 @@ pub enum WriteError {
 
 use WriteError::*;
 
+impl VirtualMem {
+    /// Allocate a VirtualMem insntace with a requested size
+    pub fn alloc(exec_mem_bytes: usize, mem_bytes: Option<usize>) -> Self {
+        let virt_block: *mut u8 = unsafe { rb_jit_reserve_addr_space(exec_mem_bytes as u32) };
+
+        // Memory protection syscalls need page-aligned addresses, so check it here. Assuming
+        // `virt_block` is page-aligned, `second_half` should be page-aligned as long as the
+        // page size in bytes is a power of two 2¹⁹ or smaller. This is because the user
+        // requested size is half of mem_option × 2²⁰ as it's in MiB.
+        //
+        // Basically, we don't support x86-64 2MiB and 1GiB pages. ARMv8 can do up to 64KiB
+        // (2¹⁶ bytes) pages, which should be fine. 4KiB pages seem to be the most popular though.
+        let page_size = unsafe { rb_jit_get_page_size() };
+        assert_eq!(
+            virt_block as usize % page_size as usize, 0,
+            "Start of virtual address block should be page-aligned",
+        );
+
+        Self::new(sys::SystemAllocator {}, page_size, NonNull::new(virt_block).unwrap(), exec_mem_bytes, mem_bytes)
+    }
+}
+
 impl<A: Allocator> VirtualMemory<A> {
     /// Bring a part of the address space under management.
     pub fn new(
         allocator: A,
         page_size: u32,
         virt_region_start: NonNull<u8>,
         region_size_bytes: usize,
-        memory_limit_bytes: usize,
+        memory_limit_bytes: Option<usize>,
     ) -> Self {
         assert_ne!(0, page_size);
         let page_size_bytes = page_size as usize;
@@ -181,6 +196,12 @@ impl<A: Allocator> VirtualMemory<A> {
             let whole_region_end = start.wrapping_add(self.region_size_bytes);
             let alloc = &mut self.allocator;
 
+            // Ignore zjit_alloc_size() if self.memory_limit_bytes is None for testing
+            let mut required_region_bytes = page_addr + page_size - start as usize;
+            if self.memory_limit_bytes.is_some() {
+                required_region_bytes += zjit_alloc_size();
+            }
+
             assert!((start..=whole_region_end).contains(&mapped_region_end));
 
             if (start..mapped_region_end).contains(&raw) {
@@ -193,7 +214,7 @@ impl<A: Allocator> VirtualMemory<A> {
 
                 self.current_write_page = Some(page_addr);
             } else if (start..whole_region_end).contains(&raw) &&
-                    (page_addr + page_size - start as usize) + zjit_alloc_size() < self.memory_limit_bytes {
+                    required_region_bytes < self.memory_limit_bytes.unwrap_or(self.region_size_bytes) {
                 // Writing to a brand new page
                 let mapped_region_end_addr = mapped_region_end as usize;
                 let alloc_size = page_addr - mapped_region_end_addr + page_size;
@@ -279,7 +300,6 @@ impl<A: Allocator> CodePtrBase for VirtualMemory<A> {
 }
 
 /// Requires linking with CRuby to work
-#[cfg(not(test))]
 pub mod sys {
     use crate::cruby::*;
 
@@ -387,7 +407,7 @@ pub mod tests {
             PAGE_SIZE.try_into().unwrap(),
             NonNull::new(mem_start as *mut u8).unwrap(),
             mem_size,
-            get_option!(mem_bytes),
+            None,
         )
     }