1 files changed, 113 insertions, 48 deletions

diff --git a/yjit/src/virtualmem.rs b/yjit/src/virtualmem.rs
index 6a8e27447e..f3c0ceefff 100644
--- a/yjit/src/virtualmem.rs
+++ b/yjit/src/virtualmem.rs
@@ -3,7 +3,9 @@
 // usize->pointer casts is viable. It seems like a lot of work for us to participate for not much
 // benefit.
 
-use crate::utils::IntoUsize;
+use std::ptr::NonNull;
+
+use crate::{utils::IntoUsize, backend::ir::Target};
 
 #[cfg(not(test))]
 pub type VirtualMem = VirtualMemory<sys::SystemAllocator>;
@@ -22,7 +24,7 @@ pub type VirtualMem = VirtualMemory<tests::TestingAllocator>;
 /// the code in the region executable.
 pub struct VirtualMemory<A: Allocator> {
     /// Location of the virtual memory region.
-    region_start: *mut u8,
+    region_start: NonNull<u8>,
 
     /// Size of the region in bytes.
     region_size_bytes: usize,
@@ -51,14 +53,47 @@ pub trait Allocator {
     fn mark_writable(&mut self, ptr: *const u8, size: u32) -> bool;
 
     fn mark_executable(&mut self, ptr: *const u8, size: u32);
+
+    fn mark_unused(&mut self, ptr: *const u8, size: u32) -> bool;
 }
 
-/// Pointer into a [VirtualMemory].
-/// We may later change this to wrap an u32.
-/// Note: there is no NULL constant for CodePtr. You should use Option<CodePtr> instead.
+/// Pointer into a [VirtualMemory] represented as an offset from the base.
+/// Note: there is no NULL constant for [CodePtr]. You should use `Option<CodePtr>` instead.
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Debug)]
-#[repr(C)]
-pub struct CodePtr(*const u8);
+#[repr(C, packed)]
+pub struct CodePtr(u32);
+
+impl CodePtr {
+    /// Advance the CodePtr. Can return a dangling pointer.
+    pub fn add_bytes(self, bytes: usize) -> Self {
+        let CodePtr(raw) = self;
+        let bytes: u32 = bytes.try_into().unwrap();
+        CodePtr(raw + bytes)
+    }
+
+    /// Note that the raw pointer might be dangling if there hasn't
+    /// been any writes to it through the [VirtualMemory] yet.
+    pub fn raw_ptr(self, base: &impl CodePtrBase) -> *const u8 {
+        let CodePtr(offset) = self;
+        return base.base_ptr().as_ptr().wrapping_add(offset.as_usize())
+    }
+
+    /// Get the address of the code pointer.
+    pub fn raw_addr(self, base: &impl CodePtrBase) -> usize {
+        self.raw_ptr(base) as usize
+    }
+
+    /// Get the offset component for the code pointer. Useful finding the distance between two
+    /// code pointers that share the same [VirtualMem].
+    pub fn as_offset(self) -> i64 {
+        let CodePtr(offset) = self;
+        offset.into()
+    }
+
+    pub fn as_side_exit(self) -> Target {
+        Target::SideExitPtr(self)
+    }
+}
 
 /// Errors that can happen when writing to [VirtualMemory]
 #[derive(Debug, PartialEq)]
@@ -71,7 +106,7 @@ use WriteError::*;
 
 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: *mut u8, size_bytes: usize) -> Self {
+    pub fn new(allocator: A, page_size: u32, virt_region_start: NonNull<u8>, size_bytes: usize) -> Self {
         assert_ne!(0, page_size);
         let page_size_bytes = page_size.as_usize();
 
@@ -88,7 +123,20 @@ impl<A: Allocator> VirtualMemory<A> {
     /// Return the start of the region as a raw pointer. Note that it could be a dangling
     /// pointer so be careful dereferencing it.
     pub fn start_ptr(&self) -> CodePtr {
-        CodePtr(self.region_start)
+        CodePtr(0)
+    }
+
+    pub fn mapped_end_ptr(&self) -> CodePtr {
+        self.start_ptr().add_bytes(self.mapped_region_bytes)
+    }
+
+    pub fn virtual_end_ptr(&self) -> CodePtr {
+        self.start_ptr().add_bytes(self.region_size_bytes)
+    }
+
+    /// Size of the region in bytes that we have allocated physical memory for.
+    pub fn mapped_region_size(&self) -> usize {
+        self.mapped_region_bytes
     }
 
     /// Size of the region in bytes where writes could be attempted.
@@ -96,17 +144,23 @@ impl<A: Allocator> VirtualMemory<A> {
         self.region_size_bytes
     }
 
+    /// The granularity at which we can control memory permission.
+    /// On Linux, this is the page size that mmap(2) talks about.
+    pub fn system_page_size(&self) -> usize {
+        self.page_size_bytes
+    }
+
     /// Write a single byte. The first write to a page makes it readable.
     pub fn write_byte(&mut self, write_ptr: CodePtr, byte: u8) -> Result<(), WriteError> {
         let page_size = self.page_size_bytes;
-        let raw: *mut u8 = write_ptr.raw_ptr() as *mut u8;
+        let raw: *mut u8 = write_ptr.raw_ptr(self) as *mut u8;
         let page_addr = (raw as usize / page_size) * page_size;
 
         if self.current_write_page == Some(page_addr) {
             // Writing within the last written to page, nothing to do
         } else {
             // Switching to a different and potentially new page
-            let start = self.region_start;
+            let start = self.region_start.as_ptr();
             let mapped_region_end = start.wrapping_add(self.mapped_region_bytes);
             let whole_region_end = start.wrapping_add(self.region_size_bytes);
             let alloc = &mut self.allocator;
@@ -141,10 +195,16 @@ impl<A: Allocator> VirtualMemory<A> {
                     if !alloc.mark_writable(mapped_region_end.cast(), alloc_size_u32) {
                         return Err(FailedPageMapping);
                     }
-                    // Fill new memory with PUSH DS (0x1E) so that executing uninitialized memory
-                    // will fault with #UD in 64-bit mode. On Linux it becomes SIGILL and use the
-                    // usual Ruby crash reporter.
-                    std::slice::from_raw_parts_mut(mapped_region_end, alloc_size).fill(0x1E);
+                    if cfg!(target_arch = "x86_64") {
+                        // Fill new memory with PUSH DS (0x1E) so that executing uninitialized memory
+                        // will fault with #UD in 64-bit mode. On Linux it becomes SIGILL and use the
+                        // usual Ruby crash reporter.
+                        std::slice::from_raw_parts_mut(mapped_region_end, alloc_size).fill(0x1E);
+                    } else if cfg!(target_arch = "aarch64") {
+                        // In aarch64, all zeros encodes UDF, so it's already what we want.
+                    } else {
+                        unreachable!("unknown arch");
+                    }
                 }
                 self.mapped_region_bytes = self.mapped_region_bytes + alloc_size;
 
@@ -169,39 +229,35 @@ impl<A: Allocator> VirtualMemory<A> {
         let mapped_region_bytes: u32 = self.mapped_region_bytes.try_into().unwrap();
 
         // Make mapped region executable
-        self.allocator.mark_executable(region_start, mapped_region_bytes);
-    }
-}
-
-impl CodePtr {
-    /// Note that the raw pointer might be dangling if there hasn't
-    /// been any writes to it through the [VirtualMemory] yet.
-    pub fn raw_ptr(self) -> *const u8 {
-        let CodePtr(ptr) = self;
-        return ptr;
-    }
-
-    /// Advance the CodePtr. Can return a dangling pointer.
-    pub fn add_bytes(self, bytes: usize) -> Self {
-        let CodePtr(raw) = self;
-        CodePtr(raw.wrapping_add(bytes))
+        self.allocator.mark_executable(region_start.as_ptr(), mapped_region_bytes);
     }
 
-    pub fn into_i64(self) -> i64 {
-        let CodePtr(ptr) = self;
-        ptr as i64
+    /// Free a range of bytes. start_ptr must be memory page-aligned.
+    pub fn free_bytes(&mut self, start_ptr: CodePtr, size: u32) {
+        assert_eq!(start_ptr.raw_ptr(self) as usize % self.page_size_bytes, 0);
+
+        // Bounds check the request. We should only free memory we manage.
+        let mapped_region = self.start_ptr().raw_ptr(self)..self.mapped_end_ptr().raw_ptr(self);
+        let virtual_region = self.start_ptr().raw_ptr(self)..self.virtual_end_ptr().raw_ptr(self);
+        let last_byte_to_free = start_ptr.add_bytes(size.saturating_sub(1).as_usize()).raw_ptr(self);
+        assert!(mapped_region.contains(&start_ptr.raw_ptr(self)));
+        // On platforms where code page size != memory page size (e.g. Linux), we often need
+        // to free code pages that contain unmapped memory pages. When it happens on the last
+        // code page, it's more appropriate to check the last byte against the virtual region.
+        assert!(virtual_region.contains(&last_byte_to_free));
+
+        self.allocator.mark_unused(start_ptr.raw_ptr(self), size);
     }
+}
 
-    pub fn into_usize(self) -> usize {
-        let CodePtr(ptr) = self;
-        ptr as usize
-    }
+/// Something that could provide a base pointer to compute a raw pointer from a [CodePtr].
+pub trait CodePtrBase {
+    fn base_ptr(&self) -> NonNull<u8>;
 }
 
-impl From<*mut u8> for CodePtr {
-    fn from(value: *mut u8) -> Self {
-        assert!(value as usize != 0);
-        return CodePtr(value);
+impl<A: Allocator> CodePtrBase for VirtualMemory<A> {
+    fn base_ptr(&self) -> NonNull<u8> {
+        self.region_start
     }
 }
 
@@ -223,6 +279,10 @@ mod sys {
         fn mark_executable(&mut self, ptr: *const u8, size: u32) {
             unsafe { rb_yjit_mark_executable(ptr as VoidPtr, size) }
         }
+
+        fn mark_unused(&mut self, ptr: *const u8, size: u32) -> bool {
+            unsafe { rb_yjit_mark_unused(ptr as VoidPtr, size) }
+        }
     }
 }
 
@@ -246,6 +306,7 @@ pub mod tests {
     enum AllocRequest {
         MarkWritable{ start_idx: usize, length: usize },
         MarkExecutable{ start_idx: usize, length: usize },
+        MarkUnused,
     }
     use AllocRequest::*;
 
@@ -286,6 +347,13 @@ pub mod tests {
             // We don't try to execute generated code in cfg(test)
             // so no need to actually request executable memory.
         }
+
+        fn mark_unused(&mut self, ptr: *const u8, length: u32) -> bool {
+            self.bounds_check_request(ptr, length);
+            self.requests.push(MarkUnused);
+
+            true
+        }
     }
 
     // Fictional architecture where each page is 4 bytes long
@@ -298,12 +366,13 @@ pub mod tests {
         VirtualMemory::new(
             alloc,
             PAGE_SIZE.try_into().unwrap(),
-            mem_start as *mut u8,
+            NonNull::new(mem_start as *mut u8).unwrap(),
             mem_size,
         )
     }
 
     #[test]
+    #[cfg(target_arch = "x86_64")]
     fn new_memory_is_initialized() {
         let mut virt = new_dummy_virt_mem();
 
@@ -340,16 +409,12 @@ pub mod tests {
     #[test]
     fn bounds_checking() {
         use super::WriteError::*;
-        use std::ptr;
         let mut virt = new_dummy_virt_mem();
 
-        let null = CodePtr(ptr::null());
-        assert_eq!(Err(OutOfBounds), virt.write_byte(null, 0));
-
         let one_past_end = virt.start_ptr().add_bytes(virt.virtual_region_size());
         assert_eq!(Err(OutOfBounds), virt.write_byte(one_past_end, 0));
 
-        let end_of_addr_space = CodePtr(usize::MAX as _);
+        let end_of_addr_space = CodePtr(u32::MAX);
         assert_eq!(Err(OutOfBounds), virt.write_byte(end_of_addr_space, 0));
     }