Copy VirtualMem and mmap

Merged: https://github.com/ruby/ruby/pull/13131

zjit.c

@@ -31,6 +31,94 @@
 
 #include <errno.h>
 
+// Address space reservation. Memory pages are mapped on an as needed basis.
+// See the Rust mm module for details.
+uint8_t *
+rb_zjit_reserve_addr_space(uint32_t mem_size)
+{
+#ifndef _WIN32
+    uint8_t *mem_block;
+
+    // On Linux
+    #if defined(MAP_FIXED_NOREPLACE) && defined(_SC_PAGESIZE)
+        uint32_t const page_size = (uint32_t)sysconf(_SC_PAGESIZE);
+        uint8_t *const cfunc_sample_addr = (void *)(uintptr_t)&rb_zjit_reserve_addr_space;
+        uint8_t *const probe_region_end = cfunc_sample_addr + INT32_MAX;
+        // Align the requested address to page size
+        uint8_t *req_addr = align_ptr(cfunc_sample_addr, page_size);
+
+        // Probe for addresses close to this function using MAP_FIXED_NOREPLACE
+        // to improve odds of being in range for 32-bit relative call instructions.
+        do {
+            mem_block = mmap(
+                req_addr,
+                mem_size,
+                PROT_NONE,
+                MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED_NOREPLACE,
+                -1,
+                0
+            );
+
+            // If we succeeded, stop
+            if (mem_block != MAP_FAILED) {
+                ruby_annotate_mmap(mem_block, mem_size, "Ruby:rb_zjit_reserve_addr_space");
+                break;
+            }
+
+            // -4MiB. Downwards to probe away from the heap. (On x86/A64 Linux
+            // main_code_addr < heap_addr, and in case we are in a shared
+            // library mapped higher than the heap, downwards is still better
+            // since it's towards the end of the heap rather than the stack.)
+            req_addr -= 4 * 1024 * 1024;
+        } while (req_addr < probe_region_end);
+
+    // On MacOS and other platforms
+    #else
+        // Try to map a chunk of memory as executable
+        mem_block = mmap(
+            (void *)rb_zjit_reserve_addr_space,
+            mem_size,
+            PROT_NONE,
+            MAP_PRIVATE | MAP_ANONYMOUS,
+            -1,
+            0
+        );
+    #endif
+
+    // Fallback
+    if (mem_block == MAP_FAILED) {
+        // Try again without the address hint (e.g., valgrind)
+        mem_block = mmap(
+            NULL,
+            mem_size,
+            PROT_NONE,
+            MAP_PRIVATE | MAP_ANONYMOUS,
+            -1,
+            0
+        );
+
+        if (mem_block != MAP_FAILED) {
+            ruby_annotate_mmap(mem_block, mem_size, "Ruby:rb_zjit_reserve_addr_space:fallback");
+        }
+    }
+
+    // Check that the memory mapping was successful
+    if (mem_block == MAP_FAILED) {
+        perror("ruby: zjit: mmap:");
+        if(errno == ENOMEM) {
+            // No crash report if it's only insufficient memory
+            exit(EXIT_FAILURE);
+        }
+        rb_bug("mmap failed");
+    }
+
+    return mem_block;
+#else
+    // Windows not supported for now
+    return NULL;
+#endif
+}
+
 void
 rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception)
 {

zjit/src/codegen.rs

@@ -0,0 +1,19 @@
+use std::rc::Rc;
+use std::cell::RefCell;
+
+use crate::virtualmem::VirtualMem;
+
+/// Block of memory into which instructions can be assembled
+pub struct CodeBlock {
+    // Memory for storing the encoded instructions
+    mem_block: Rc<RefCell<VirtualMem>>,
+}
+
+/// Global state needed for code generation
+pub struct ZJITState {
+    /// Inline code block (fast path)
+    inline_cb: CodeBlock,
+}
+
+/// Private singleton instance of the codegen globals
+static mut ZJIT_STATE: Option<ZJITState> = None;

zjit/src/cruby_bindings.inc.rs

@@ -1154,6 +1154,10 @@ extern "C" {
     ) -> VALUE;
     pub fn rb_yjit_get_page_size() -> u32;
     pub fn rb_yjit_reserve_addr_space(mem_size: u32) -> *mut u8;
+    pub fn rb_zjit_mark_writable(mem_block: *mut ::std::os::raw::c_void, mem_size: u32) -> bool;
+    pub fn rb_zjit_mark_executable(mem_block: *mut ::std::os::raw::c_void, mem_size: u32);
+    pub fn rb_zjit_mark_unused(mem_block: *mut ::std::os::raw::c_void, mem_size: u32) -> bool;
+    pub fn rb_zjit_reserve_addr_space(mem_size: u32) -> *mut u8;
     pub fn rb_c_method_tracing_currently_enabled(ec: *const rb_execution_context_t) -> bool;
     pub fn rb_full_cfunc_return(ec: *mut rb_execution_context_t, return_value: VALUE);
     pub fn rb_iseq_encoded_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;

zjit/src/lib.rs

@@ -1,8 +1,11 @@
 #![allow(dead_code)]
 
+mod codegen;
 mod cruby;
-mod stats;
 mod ir;
+mod stats;
+mod utils;
+mod virtualmem;
 use crate::cruby::*;
 
 #[allow(non_upper_case_globals)]

zjit/src/stats.rs

@@ -3,3 +3,7 @@
 // We could also tag which stats are fallback or exit counters, etc. Maybe even tag units?
 //
 // Comptime vs Runtime stats?
+
+pub fn zjit_alloc_size() -> usize {
+    0 // TODO: report the actual memory usage
+}

zjit/src/utils.rs

@@ -0,0 +1,46 @@
+/// Trait for casting to [usize] that allows you to say `.as_usize()`.
+/// Implementation conditional on the cast preserving the numeric value on
+/// all inputs and being inexpensive.
+///
+/// [usize] is only guaranteed to be more than 16-bit wide, so we can't use
+/// `.into()` to cast an `u32` or an `u64` to a `usize` even though in all
+/// the platforms YJIT supports these two casts are pretty much no-ops.
+/// We could say `as usize` or `.try_convert().unwrap()` everywhere
+/// for those casts but they both have undesirable consequences if and when
+/// we decide to support 32-bit platforms. Unfortunately we can't implement
+/// [::core::convert::From] for [usize] since both the trait and the type are
+/// external. Naming the method `into()` also runs into naming conflicts.
+pub(crate) trait IntoUsize {
+    /// Convert to usize. Implementation conditional on width of [usize].
+    fn as_usize(self) -> usize;
+}
+
+#[cfg(target_pointer_width = "64")]
+impl IntoUsize for u64 {
+    fn as_usize(self) -> usize {
+        self as usize
+    }
+}
+
+#[cfg(target_pointer_width = "64")]
+impl IntoUsize for u32 {
+    fn as_usize(self) -> usize {
+        self as usize
+    }
+}
+
+impl IntoUsize for u16 {
+    /// Alias for `.into()`. For convenience so you could use the trait for
+    /// all unsgined types.
+    fn as_usize(self) -> usize {
+        self.into()
+    }
+}
+
+impl IntoUsize for u8 {
+    /// Alias for `.into()`. For convenience so you could use the trait for
+    /// all unsgined types.
+    fn as_usize(self) -> usize {
+        self.into()
+    }
+}

zjit/src/virtualmem.rs

@@ -0,0 +1,445 @@
+//! Memory management stuff for ZJIT's code storage. Deals with virtual memory.
+// I'm aware that there is an experiment in Rust Nightly right now for to see if banning
+// usize->pointer casts is viable. It seems like a lot of work for us to participate for not much
+// benefit.
+
+use std::ptr::NonNull;
+
+use crate::{stats::zjit_alloc_size, utils::IntoUsize};
+
+#[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
+/// `mprotect` with `PROT_READ|PROT_WRITE` as needed. The WIN32 equivalent seems to be
+/// `VirtualAlloc` with `MEM_RESERVE` then later with `MEM_COMMIT`.
+///
+/// This handles ["W^X"](https://en.wikipedia.org/wiki/W%5EX) semi-automatically. Writes
+/// are always accepted and once writes are done a call to [Self::mark_all_executable] makes
+/// the code in the region executable.
+pub struct VirtualMemory<A: Allocator> {
+    /// Location of the virtual memory region.
+    region_start: NonNull<u8>,
+
+    /// Size of this virtual memory region in bytes.
+    region_size_bytes: usize,
+
+    /// mapped_region_bytes + zjit_alloc_size may not increase beyond this limit.
+    memory_limit_bytes: usize,
+
+    /// Number of bytes per "page", memory protection permission can only be controlled at this
+    /// granularity.
+    page_size_bytes: usize,
+
+    /// Number of bytes that have we have allocated physical memory for starting at
+    /// [Self::region_start].
+    mapped_region_bytes: usize,
+
+    /// Keep track of the address of the last written to page.
+    /// Used for changing protection to implement W^X.
+    current_write_page: Option<usize>,
+
+    /// Zero size member for making syscalls to get physical memory during normal operation.
+    /// When testing this owns some memory.
+    allocator: A,
+}
+
+/// Groups together the two syscalls to get get new physical memory and to change
+/// memory protection. See [VirtualMemory] for details.
+pub trait Allocator {
+    #[must_use]
+    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] 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, 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()
+    }
+}
+
+/// Errors that can happen when writing to [VirtualMemory]
+#[derive(Debug, PartialEq)]
+pub enum WriteError {
+    OutOfBounds,
+    FailedPageMapping,
+}
+
+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: NonNull<u8>,
+        region_size_bytes: usize,
+        memory_limit_bytes: usize,
+    ) -> Self {
+        assert_ne!(0, page_size);
+        let page_size_bytes = page_size.as_usize();
+
+        Self {
+            region_start: virt_region_start,
+            region_size_bytes,
+            memory_limit_bytes,
+            page_size_bytes,
+            mapped_region_bytes: 0,
+            current_write_page: None,
+            allocator,
+        }
+    }
+
+    /// 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(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.
+    pub fn virtual_region_size(&self) -> usize {
+        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(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.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;
+
+            assert!((start..=whole_region_end).contains(&mapped_region_end));
+
+            if (start..mapped_region_end).contains(&raw) {
+                // Writing to a previously written to page.
+                // Need to make page writable, but no need to fill.
+                let page_size: u32 = page_size.try_into().unwrap();
+                if !alloc.mark_writable(page_addr as *const _, page_size) {
+                    return Err(FailedPageMapping);
+                }
+
+                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 {
+                // 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;
+
+                assert_eq!(0, alloc_size % page_size, "allocation size should be page aligned");
+                assert_eq!(0, mapped_region_end_addr % page_size, "pointer should be page aligned");
+
+                if alloc_size > page_size {
+                    // This is unusual for the current setup, so keep track of it.
+                    //crate::stats::incr_counter!(exec_mem_non_bump_alloc); // TODO
+                }
+
+                // Allocate new chunk
+                let alloc_size_u32: u32 = alloc_size.try_into().unwrap();
+                unsafe {
+                    if !alloc.mark_writable(mapped_region_end.cast(), alloc_size_u32) {
+                        return Err(FailedPageMapping);
+                    }
+                    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;
+
+                self.current_write_page = Some(page_addr);
+            } else {
+                return Err(OutOfBounds);
+            }
+        }
+
+        // We have permission to write if we get here
+        unsafe { raw.write(byte) };
+
+        Ok(())
+    }
+
+    /// 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) {
+        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 executable
+        self.allocator.mark_executable(region_start.as_ptr(), mapped_region_bytes);
+    }
+
+    /// 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);
+    }
+}
+
+/// 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<A: Allocator> CodePtrBase for VirtualMemory<A> {
+    fn base_ptr(&self) -> NonNull<u8> {
+        self.region_start
+    }
+}
+
+/// Requires linking with CRuby to work
+#[cfg(not(test))]
+mod sys {
+    use crate::cruby::*;
+
+    /// Zero size! This just groups together syscalls that require linking with CRuby.
+    pub struct SystemAllocator;
+
+    type VoidPtr = *mut std::os::raw::c_void;
+
+    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) }
+        }
+
+        fn mark_executable(&mut self, ptr: *const u8, size: u32) {
+            unsafe { rb_zjit_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) }
+        }
+    }
+}
+
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+
+    // Track allocation requests and owns some fixed size backing memory for requests.
+    // While testing we don't execute generated code.
+    pub struct TestingAllocator {
+        requests: Vec<AllocRequest>,
+        memory: Vec<u8>,
+    }
+
+    #[derive(Debug)]
+    enum AllocRequest {
+        MarkWritable{ start_idx: usize, length: usize },
+        MarkExecutable{ start_idx: usize, length: usize },
+        MarkUnused,
+    }
+    use AllocRequest::*;
+
+    impl TestingAllocator {
+        pub fn new(mem_size: usize) -> Self {
+            Self { requests: Vec::default(), memory: vec![0; mem_size] }
+        }
+
+        pub fn mem_start(&self) -> *const u8 {
+            self.memory.as_ptr()
+        }
+
+        // Verify that write_byte() bounds checks. Return `ptr` as an index.
+        fn bounds_check_request(&self, ptr: *const u8, size: u32) -> usize {
+            let mem_start = self.memory.as_ptr() as usize;
+            let index = ptr as usize - mem_start;
+
+            assert!(index < self.memory.len());
+            assert!(index + size.as_usize() <= self.memory.len());
+
+            index
+        }
+    }
+
+    // Bounds check and then record the request
+    impl super::Allocator for TestingAllocator {
+        fn mark_writable(&mut self, ptr: *const u8, length: u32) -> bool {
+            let index = self.bounds_check_request(ptr, length);
+            self.requests.push(MarkWritable { start_idx: index, length: length.as_usize() });
+
+            true
+        }
+
+        fn mark_executable(&mut self, ptr: *const u8, length: u32) {
+            let index = self.bounds_check_request(ptr, length);
+            self.requests.push(MarkExecutable { start_idx: index, length: length.as_usize() });
+
+            // 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
+    const PAGE_SIZE: usize = 4;
+    fn new_dummy_virt_mem() -> VirtualMemory<TestingAllocator> {
+        let mem_size = PAGE_SIZE * 10;
+        let alloc = TestingAllocator::new(mem_size);
+        let mem_start: *const u8 = alloc.mem_start();
+
+        VirtualMemory::new(
+            alloc,
+            PAGE_SIZE.try_into().unwrap(),
+            NonNull::new(mem_start as *mut u8).unwrap(),
+            mem_size,
+            128 * 1024 * 1024,
+        )
+    }
+
+    #[test]
+    #[cfg(target_arch = "x86_64")]
+    fn new_memory_is_initialized() {
+        let mut virt = new_dummy_virt_mem();
+
+        virt.write_byte(virt.start_ptr(), 1).unwrap();
+        assert!(
+            virt.allocator.memory[..PAGE_SIZE].iter().all(|&byte| byte != 0),
+            "Entire page should be initialized",
+        );
+
+        // Skip a few page
+        let three_pages = 3 * PAGE_SIZE;
+        virt.write_byte(virt.start_ptr().add_bytes(three_pages), 1).unwrap();
+        assert!(
+            virt.allocator.memory[..three_pages].iter().all(|&byte| byte != 0),
+            "Gaps between write requests should be filled",
+        );
+    }
+
+    #[test]
+    fn no_redundant_syscalls_when_writing_to_the_same_page() {
+        let mut virt = new_dummy_virt_mem();
+
+        virt.write_byte(virt.start_ptr(), 1).unwrap();
+        virt.write_byte(virt.start_ptr(), 0).unwrap();
+
+        assert!(
+            matches!(
+                virt.allocator.requests[..],
+                [MarkWritable { start_idx: 0, length: PAGE_SIZE }],
+            )
+        );
+    }
+
+    #[test]
+    fn bounds_checking() {
+        use super::WriteError::*;
+        let mut virt = new_dummy_virt_mem();
+
+        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(u32::MAX);
+        assert_eq!(Err(OutOfBounds), virt.write_byte(end_of_addr_space, 0));
+    }
+
+    #[test]
+    fn only_written_to_regions_become_executable() {
+        // ... so we catch attempts to read/write/execute never-written-to regions
+        const THREE_PAGES: usize = PAGE_SIZE * 3;
+        let mut virt = new_dummy_virt_mem();
+        let page_two_start = virt.start_ptr().add_bytes(PAGE_SIZE * 2);
+        virt.write_byte(page_two_start, 1).unwrap();
+        virt.mark_all_executable();
+
+        assert!(virt.virtual_region_size() > THREE_PAGES);
+        assert!(
+            matches!(
+                virt.allocator.requests[..],
+                [
+                    MarkWritable { start_idx: 0, length: THREE_PAGES },
+                    MarkExecutable { start_idx: 0, length: THREE_PAGES },
+                ]
+            ),
+        );
+    }
+}