Rust YJIT

In December 2021, we opened an [issue] to solicit feedback regarding the
porting of the YJIT codebase from C99 to Rust. There were some
reservations, but this project was given the go ahead by Ruby core
developers and Matz. Since then, we have successfully completed the port
of YJIT to Rust.

The new Rust version of YJIT has reached parity with the C version, in
that it passes all the CRuby tests, is able to run all of the YJIT
benchmarks, and performs similarly to the C version (because it works
the same way and largely generates the same machine code). We've even
incorporated some design improvements, such as a more fine-grained
constant invalidation mechanism which we expect will make a big
difference in Ruby on Rails applications.

Because we want to be careful, YJIT is guarded behind a configure
option:

```shell
./configure --enable-yjit # Build YJIT in release mode
./configure --enable-yjit=dev # Build YJIT in dev/debug mode
```

By default, YJIT does not get compiled and cargo/rustc is not required.
If YJIT is built in dev mode, then `cargo` is used to fetch development
dependencies, but when building in release, `cargo` is not required,
only `rustc`. At the moment YJIT requires Rust 1.60.0 or newer.

The YJIT command-line options remain mostly unchanged, and more details
about the build process are documented in `doc/yjit/yjit.md`.

The CI tests have been updated and do not take any more resources than
before.

The development history of the Rust port is available at the following
commit for interested parties:
https://github.com/Shopify/ruby/commit/1fd9573d8b4b65219f1c2407f30a0a60e537f8be

Our hope is that Rust YJIT will be compiled and included as a part of
system packages and compiled binaries of the Ruby 3.2 release. We do not
anticipate any major problems as Rust is well supported on every
platform which YJIT supports, but to make sure that this process works
smoothly, we would like to reach out to those who take care of building
systems packages before the 3.2 release is shipped and resolve any
issues that may come up.

[issue]: https://bugs.ruby-lang.org/issues/18481

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>
Co-authored-by: Noah Gibbs <the.codefolio.guy@gmail.com>
Co-authored-by: Kevin Newton <kddnewton@gmail.com>

1 files changed, 392 insertions, 0 deletions

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
new file mode 100644
index 0000000000..0d61cd654a
--- /dev/null
+++ b/yjit/src/asm/mod.rs
@@ -0,0 +1,392 @@
+use std::collections::BTreeMap;
+use std::mem;
+
+// Lots of manual vertical alignment in there that rustfmt doesn't handle well.
+#[rustfmt::skip]
+pub mod x86_64;
+
+/// Pointer to a piece of machine code
+/// We may later change this to wrap an u32
+/// 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);
+
+impl CodePtr {
+    pub fn raw_ptr(&self) -> *const u8 {
+        let CodePtr(ptr) = *self;
+        return ptr;
+    }
+
+    fn into_i64(&self) -> i64 {
+        let CodePtr(ptr) = self;
+        *ptr as i64
+    }
+
+    fn into_usize(&self) -> usize {
+        let CodePtr(ptr) = self;
+        *ptr as usize
+    }
+}
+
+impl From<*mut u8> for CodePtr {
+    fn from(value: *mut u8) -> Self {
+        assert!(value as usize != 0);
+        return CodePtr(value);
+    }
+}
+
+/// Compute an offset in bytes of a given struct field
+macro_rules! offset_of {
+    ($struct_type:ty, $field_name:tt) => {{
+        // Null pointer to our struct type
+        let foo = (0 as *const $struct_type);
+
+        unsafe {
+            let ptr_field = (&(*foo).$field_name as *const _ as usize);
+            let ptr_base = (foo as usize);
+            ptr_field - ptr_base
+        }
+    }};
+}
+pub(crate) use offset_of;
+
+//
+// TODO: need a field_size_of macro, to compute the size of a struct field in bytes
+//
+
+// 1 is not aligned so this won't match any pages
+const ALIGNED_WRITE_POSITION_NONE: usize = 1;
+
+/// Reference to an ASM label
+struct LabelRef {
+    // Position in the code block where the label reference exists
+    pos: usize,
+
+    // Label which this refers to
+    label_idx: usize,
+}
+
+/// Block of memory into which instructions can be assembled
+pub struct CodeBlock {
+    // Block of non-executable memory used for dummy code blocks
+    // This memory is owned by this block and lives as long as the block
+    dummy_block: Vec<u8>,
+
+    // Pointer to memory we are writing into
+    mem_block: *mut u8,
+
+    // Memory block size
+    mem_size: usize,
+
+    // Current writing position
+    write_pos: usize,
+
+    // Table of registered label addresses
+    label_addrs: Vec<usize>,
+
+    // Table of registered label names
+    label_names: Vec<String>,
+
+    // References to labels
+    label_refs: Vec<LabelRef>,
+
+    // Comments for assembly instructions, if that feature is enabled
+    asm_comments: BTreeMap<usize, Vec<String>>,
+
+    // Keep track of the current aligned write position.
+    // Used for changing protection when writing to the JIT buffer
+    current_aligned_write_pos: usize,
+
+    // Memory protection works at page granularity and this is the
+    // the size of each page. Used to implement W^X.
+    page_size: usize,
+
+    // Set if the CodeBlock is unable to output some instructions,
+    // for example, when there is not enough space or when a jump
+    // target is too far away.
+    dropped_bytes: bool,
+}
+
+impl CodeBlock {
+    pub fn new_dummy(mem_size: usize) -> Self {
+        // Allocate some non-executable memory
+        let mut dummy_block = vec![0; mem_size];
+        let mem_ptr = dummy_block.as_mut_ptr();
+
+        Self {
+            dummy_block: dummy_block,
+            mem_block: mem_ptr,
+            mem_size: mem_size,
+            write_pos: 0,
+            label_addrs: Vec::new(),
+            label_names: Vec::new(),
+            label_refs: Vec::new(),
+            asm_comments: BTreeMap::new(),
+            current_aligned_write_pos: ALIGNED_WRITE_POSITION_NONE,
+            page_size: 4096,
+            dropped_bytes: false,
+        }
+    }
+
+    pub fn new(mem_block: *mut u8, mem_size: usize, page_size: usize) -> Self {
+        Self {
+            dummy_block: vec![0; 0],
+            mem_block: mem_block,
+            mem_size: mem_size,
+            write_pos: 0,
+            label_addrs: Vec::new(),
+            label_names: Vec::new(),
+            label_refs: Vec::new(),
+            asm_comments: BTreeMap::new(),
+            current_aligned_write_pos: ALIGNED_WRITE_POSITION_NONE,
+            page_size,
+            dropped_bytes: false,
+        }
+    }
+
+    // Check if this code block has sufficient remaining capacity
+    pub fn has_capacity(&self, num_bytes: usize) -> bool {
+        self.write_pos + num_bytes < self.mem_size
+    }
+
+    /// Add an assembly comment if the feature is on.
+    /// If not, this becomes an inline no-op.
+    #[inline]
+    pub fn add_comment(&mut self, comment: &str) {
+        if cfg!(feature = "asm_comments") {
+            let cur_ptr = self.get_write_ptr().into_usize();
+            let this_line_comments = self.asm_comments.get(&cur_ptr);
+
+            // If there's no current list of comments for this line number, add one.
+            if this_line_comments.is_none() {
+                let new_comments = Vec::new();
+                self.asm_comments.insert(cur_ptr, new_comments);
+            }
+            let this_line_comments = self.asm_comments.get_mut(&cur_ptr).unwrap();
+
+            // Unless this comment is the same as the last one at this same line, add it.
+            let string_comment = String::from(comment);
+            if this_line_comments.last() != Some(&string_comment) {
+                this_line_comments.push(string_comment);
+            }
+        }
+    }
+
+    pub fn comments_at(&self, pos: usize) -> Option<&Vec<String>> {
+        self.asm_comments.get(&pos)
+    }
+
+    pub fn get_mem_size(&self) -> usize {
+        self.mem_size
+    }
+
+    pub fn get_write_pos(&self) -> usize {
+        self.write_pos
+    }
+
+    // Set the current write position
+    pub fn set_pos(&mut self, pos: usize) {
+        // Assert here since while CodeBlock functions do bounds checking, there is
+        // nothing stopping users from taking out an out-of-bounds pointer and
+        // doing bad accesses with it.
+        assert!(pos < self.mem_size);
+        self.write_pos = pos;
+    }
+
+    // Align the current write pointer to a multiple of bytes
+    pub fn align_pos(&mut self, multiple: u32) {
+        // Compute the alignment boundary that is lower or equal
+        // Do everything with usize
+        let multiple: usize = multiple.try_into().unwrap();
+        let pos = self.get_write_ptr().raw_ptr() as usize;
+        let remainder = pos % multiple;
+        let prev_aligned = pos - remainder;
+
+        if prev_aligned == pos {
+            // Already aligned so do nothing
+        } else {
+            // Align by advancing
+            let pad = multiple - remainder;
+            self.set_pos(self.get_write_pos() + pad);
+        }
+    }
+
+    // Set the current write position from a pointer
+    pub fn set_write_ptr(&mut self, code_ptr: CodePtr) {
+        let pos = (code_ptr.raw_ptr() as usize) - (self.mem_block as usize);
+        self.set_pos(pos);
+    }
+
+    // Get a direct pointer into the executable memory block
+    pub fn get_ptr(&self, offset: usize) -> CodePtr {
+        unsafe {
+            let ptr = self.mem_block.offset(offset as isize);
+            CodePtr(ptr)
+        }
+    }
+
+    // Get a direct pointer to the current write position
+    pub fn get_write_ptr(&mut self) -> CodePtr {
+        self.get_ptr(self.write_pos)
+    }
+
+    // Write a single byte at the current position
+    pub fn write_byte(&mut self, byte: u8) {
+        if self.write_pos < self.mem_size {
+            self.mark_position_writable(self.write_pos);
+            unsafe { self.mem_block.add(self.write_pos).write(byte) };
+            self.write_pos += 1;
+        } else {
+            self.dropped_bytes = true;
+        }
+    }
+
+    // Read a single byte at the given position
+    pub fn read_byte(&self, pos: usize) -> u8 {
+        assert!(pos < self.mem_size);
+        unsafe { self.mem_block.add(pos).read() }
+    }
+
+    // Write multiple bytes starting from the current position
+    pub fn write_bytes(&mut self, bytes: &[u8]) {
+        for byte in bytes {
+            self.write_byte(*byte);
+        }
+    }
+
+    // Write a signed integer over a given number of bits at the current position
+    pub fn write_int(&mut self, val: u64, num_bits: u32) {
+        assert!(num_bits > 0);
+        assert!(num_bits % 8 == 0);
+
+        // Switch on the number of bits
+        match num_bits {
+            8 => self.write_byte(val as u8),
+            16 => self.write_bytes(&[(val & 0xff) as u8, ((val >> 8) & 0xff) as u8]),
+            32 => self.write_bytes(&[
+                (val & 0xff) as u8,
+                ((val >> 8) & 0xff) as u8,
+                ((val >> 16) & 0xff) as u8,
+                ((val >> 24) & 0xff) as u8,
+            ]),
+            _ => {
+                let mut cur = val;
+
+                // Write out the bytes
+                for _byte in 0..(num_bits / 8) {
+                    self.write_byte((cur & 0xff) as u8);
+                    cur >>= 8;
+                }
+            }
+        }
+    }
+
+    /// Check if bytes have been dropped (unwritten because of insufficient space)
+    pub fn has_dropped_bytes(&self) -> bool {
+        self.dropped_bytes
+    }
+
+    /// Allocate a new label with a given name
+    pub fn new_label(&mut self, name: String) -> usize {
+        // This label doesn't have an address yet
+        self.label_addrs.push(0);
+        self.label_names.push(name);
+
+        return self.label_addrs.len() - 1;
+    }
+
+    /// Write a label at the current address
+    pub fn write_label(&mut self, label_idx: usize) {
+        // TODO: make sure that label_idx is valid
+        // TODO: add an asseer here
+
+        self.label_addrs[label_idx] = self.write_pos;
+    }
+
+    // Add a label reference at the current write position
+    pub fn label_ref(&mut self, label_idx: usize) {
+        // TODO: make sure that label_idx is valid
+        // TODO: add an asseer here
+
+        // Keep track of the reference
+        self.label_refs.push(LabelRef {
+            pos: self.write_pos,
+            label_idx,
+        });
+    }
+
+    // Link internal label references
+    pub fn link_labels(&mut self) {
+        let orig_pos = self.write_pos;
+
+        // For each label reference
+        for label_ref in mem::take(&mut self.label_refs) {
+            let ref_pos = label_ref.pos;
+            let label_idx = label_ref.label_idx;
+            assert!(ref_pos < self.mem_size);
+
+            let label_addr = self.label_addrs[label_idx];
+            assert!(label_addr < self.mem_size);
+
+            // Compute the offset from the reference's end to the label
+            let offset = (label_addr as i64) - ((ref_pos + 4) as i64);
+
+            self.set_pos(ref_pos);
+            self.write_int(offset as u64, 32);
+        }
+
+        self.write_pos = orig_pos;
+
+        // Clear the label positions and references
+        self.label_addrs.clear();
+        self.label_names.clear();
+        assert!(self.label_refs.is_empty());
+    }
+
+    pub fn mark_position_writable(&mut self, write_pos: usize) {
+        let page_size = self.page_size;
+        let aligned_position = (write_pos / page_size) * page_size;
+
+        if self.current_aligned_write_pos != aligned_position {
+            self.current_aligned_write_pos = aligned_position;
+
+            #[cfg(not(test))]
+            unsafe {
+                use core::ffi::c_void;
+                let page_ptr = self.get_ptr(aligned_position).raw_ptr() as *mut c_void;
+                crate::cruby::rb_yjit_mark_writable(page_ptr, page_size.try_into().unwrap());
+            }
+        }
+    }
+
+    pub fn mark_all_executable(&mut self) {
+        self.current_aligned_write_pos = ALIGNED_WRITE_POSITION_NONE;
+
+        #[cfg(not(test))]
+        unsafe {
+            use core::ffi::c_void;
+            // NOTE(alan): Right now we do allocate one big chunck and give the top half to the outlined codeblock
+            // The start of the top half of the region isn't necessarily a page boundary...
+            let cb_start = self.get_ptr(0).raw_ptr() as *mut c_void;
+            crate::cruby::rb_yjit_mark_executable(cb_start, self.mem_size.try_into().unwrap());
+        }
+    }
+}
+
+/// Wrapper struct so we can use the type system to distinguish
+/// Between the inlined and outlined code blocks
+pub struct OutlinedCb {
+    // This must remain private
+    cb: CodeBlock,
+}
+
+impl OutlinedCb {
+    pub fn wrap(cb: CodeBlock) -> Self {
+        OutlinedCb { cb: cb }
+    }
+
+    pub fn unwrap(&mut self) -> &mut CodeBlock {
+        &mut self.cb
+    }
+}