1 files changed, 67 insertions, 15 deletions

diff --git a/zjit/src/virtualmem.rs b/zjit/src/virtualmem.rs
index 42ce525fde..0088ef1a66 100644
--- a/zjit/src/virtualmem.rs
+++ b/zjit/src/virtualmem.rs
@@ -4,15 +4,11 @@
 // benefit.
 
 use std::ptr::NonNull;
+use crate::cruby::*;
+use crate::stats::zjit_alloc_bytes;
 
-use crate::stats::zjit_alloc_size;
-
-#[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
@@ -30,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.
@@ -90,7 +86,7 @@ impl CodePtr {
 
     /// Get the address of the code pointer.
     pub fn raw_addr(self, base: &impl CodePtrBase) -> usize {
-        self.raw_ptr(base) as usize
+        self.raw_ptr(base).addr()
     }
 
     /// Get the offset component for the code pointer. Useful finding the distance between two
@@ -110,6 +106,28 @@ 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^19 or smaller. This is because the user
+        // requested size is half of mem_option * 2^20 as it's in MiB.
+        //
+        // Basically, we don't support x86-64 2MiB and 1GiB pages. ARMv8 can do up to 64KiB
+        // (2^16 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(
@@ -117,7 +135,7 @@ impl<A: Allocator> VirtualMemory<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;
@@ -178,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_bytes();
+            }
+
             assert!((start..=whole_region_end).contains(&mapped_region_end));
 
             if (start..mapped_region_end).contains(&raw) {
@@ -190,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;
@@ -234,6 +258,29 @@ impl<A: Allocator> VirtualMemory<A> {
         Ok(())
     }
 
+    /// Return true if write_byte() can allocate a new page
+    pub fn can_allocate(&self) -> bool {
+        let memory_usage_bytes = self.mapped_region_bytes + zjit_alloc_bytes();
+        let memory_limit_bytes = self.memory_limit_bytes.unwrap_or(self.region_size_bytes);
+        memory_usage_bytes + self.page_size_bytes < memory_limit_bytes
+    }
+
+    /// Make all the code in the region writable. Call this before bulk writes (e.g. GC
+    /// reference updates). See [Self] for usual usage flow.
+    pub fn mark_all_writable(&mut self) {
+        self.current_write_page = None;
+
+        let region_start = self.region_start;
+        let mapped_region_bytes: u32 = self.mapped_region_bytes.try_into().unwrap();
+
+        // Make mapped region writable
+        if mapped_region_bytes > 0 {
+            if !self.allocator.mark_writable(region_start.as_ptr(), mapped_region_bytes) {
+                panic!("Cannot make JIT memory region writable");
+            }
+        }
+    }
+
     /// Make all the code in the region executable. Call this at the end of a write session.
     /// See [Self] for usual usage flow.
     pub fn mark_all_executable(&mut self) {
@@ -276,7 +323,6 @@ impl<A: Allocator> CodePtrBase for VirtualMemory<A> {
 }
 
 /// Requires linking with CRuby to work
-#[cfg(not(test))]
 pub mod sys {
     use crate::cruby::*;
 
@@ -287,15 +333,15 @@ pub mod sys {
 
     impl super::Allocator for SystemAllocator {
         fn mark_writable(&mut self, ptr: *const u8, size: u32) -> bool {
-            unsafe { rb_zjit_mark_writable(ptr as VoidPtr, size) }
+            unsafe { rb_jit_mark_writable(ptr as VoidPtr, size) }
         }
 
         fn mark_executable(&mut self, ptr: *const u8, size: u32) {
-            unsafe { rb_zjit_mark_executable(ptr as VoidPtr, size) }
+            unsafe { rb_jit_mark_executable(ptr as VoidPtr, size) }
         }
 
         fn mark_unused(&mut self, ptr: *const u8, size: u32) -> bool {
-            unsafe { rb_zjit_mark_unused(ptr as VoidPtr, size) }
+            unsafe { rb_jit_mark_unused(ptr as VoidPtr, size) }
         }
     }
 }
@@ -369,6 +415,12 @@ pub mod tests {
     // Fictional architecture where each page is 4 bytes long
     const PAGE_SIZE: usize = 4;
     fn new_dummy_virt_mem() -> VirtualMemory<TestingAllocator> {
+        unsafe {
+            if crate::options::OPTIONS.is_none() {
+                crate::options::OPTIONS = Some(crate::options::Options::default());
+            }
+        }
+
         let mem_size = PAGE_SIZE * 10;
         let alloc = TestingAllocator::new(mem_size);
         let mem_start: *const u8 = alloc.mem_start();
@@ -378,7 +430,7 @@ pub mod tests {
             PAGE_SIZE.try_into().unwrap(),
             NonNull::new(mem_start as *mut u8).unwrap(),
             mem_size,
-            128 * 1024 * 1024,
+            None,
         )
     }