3 files changed, 35 insertions, 25 deletions

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index 0320fdd829..9ef675b34d 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -1,4 +1,3 @@
-use std::cell::RefCell;
 use std::fmt;
 use std::mem;
 use std::rc::Rc;
@@ -44,7 +43,7 @@ pub struct LabelRef {
 /// Block of memory into which instructions can be assembled
 pub struct CodeBlock {
     // Memory for storing the encoded instructions
-    mem_block: Rc<RefCell<VirtualMem>>,
+    mem_block: Rc<VirtualMem>,
 
     // Size of a code page in bytes. Each code page is split into an inlined and an outlined portion.
     // Code GC collects code memory at this granularity.
@@ -107,16 +106,16 @@ impl CodeBlock {
     const PREFERRED_CODE_PAGE_SIZE: usize = 16 * 1024;
 
     /// Make a new CodeBlock
-    pub fn new(mem_block: Rc<RefCell<VirtualMem>>, outlined: bool, freed_pages: Rc<Option<Vec<usize>>>, keep_comments: bool) -> Self {
+    pub fn new(mem_block: Rc<VirtualMem>, outlined: bool, freed_pages: Rc<Option<Vec<usize>>>, keep_comments: bool) -> Self {
         // Pick the code page size
-        let system_page_size = mem_block.borrow().system_page_size();
+        let system_page_size = mem_block.system_page_size();
         let page_size = if 0 == Self::PREFERRED_CODE_PAGE_SIZE % system_page_size {
             Self::PREFERRED_CODE_PAGE_SIZE
         } else {
             system_page_size
         };
 
-        let mem_size = mem_block.borrow().virtual_region_size();
+        let mem_size = mem_block.virtual_region_size();
         let mut cb = Self {
             mem_block,
             mem_size,
@@ -238,9 +237,9 @@ impl CodeBlock {
             }
 
             // Free the grouped pages at once
-            let start_ptr = self.mem_block.borrow().start_ptr().add_bytes(page_idx * self.page_size);
+            let start_ptr = self.mem_block.start_ptr().add_bytes(page_idx * self.page_size);
             let batch_size = self.page_size * batch_idxs.len();
-            self.mem_block.borrow_mut().free_bytes(start_ptr, batch_size as u32);
+            self.mem_block.free_bytes(start_ptr, batch_size as u32);
         }
     }
 
@@ -249,13 +248,13 @@ impl CodeBlock {
     }
 
     pub fn mapped_region_size(&self) -> usize {
-        self.mem_block.borrow().mapped_region_size()
+        self.mem_block.mapped_region_size()
     }
 
     /// Size of the region in bytes where writes could be attempted.
     #[cfg(target_arch = "aarch64")]
     pub fn virtual_region_size(&self) -> usize {
-        self.mem_block.borrow().virtual_region_size()
+        self.mem_block.virtual_region_size()
     }
 
     /// Return the number of code pages that have been mapped by the VirtualMemory.
@@ -267,7 +266,7 @@ impl CodeBlock {
 
     /// Return the number of code pages that have been reserved by the VirtualMemory.
     pub fn num_virtual_pages(&self) -> usize {
-        let virtual_region_size = self.mem_block.borrow().virtual_region_size();
+        let virtual_region_size = self.mem_block.virtual_region_size();
         // CodeBlock's page size != VirtualMem's page size on Linux,
         // so mapped_region_size % self.page_size may not be 0
         ((virtual_region_size - 1) / self.page_size) + 1
@@ -409,7 +408,7 @@ impl CodeBlock {
     }
 
     pub fn write_mem(&self, write_ptr: CodePtr, byte: u8) -> Result<(), WriteError> {
-        self.mem_block.borrow_mut().write_byte(write_ptr, byte)
+        self.mem_block.write_byte(write_ptr, byte)
     }
 
     // Set the current write position
@@ -423,31 +422,31 @@ impl CodeBlock {
 
     // Set the current write position from a pointer
     pub fn set_write_ptr(&mut self, code_ptr: CodePtr) {
-        let pos = code_ptr.as_offset() - self.mem_block.borrow().start_ptr().as_offset();
+        let pos = code_ptr.as_offset() - self.mem_block.start_ptr().as_offset();
         self.set_pos(pos.try_into().unwrap());
     }
 
     /// Get a (possibly dangling) direct pointer into the executable memory block
     pub fn get_ptr(&self, offset: usize) -> CodePtr {
-        self.mem_block.borrow().start_ptr().add_bytes(offset)
+        self.mem_block.start_ptr().add_bytes(offset)
     }
 
     /// Convert an address range to memory page indexes against a num_pages()-sized array.
     pub fn addrs_to_pages(&self, start_addr: CodePtr, end_addr: CodePtr) -> impl Iterator<Item = usize> {
-        let mem_start = self.mem_block.borrow().start_ptr().raw_addr(self);
-        let mem_end = self.mem_block.borrow().mapped_end_ptr().raw_addr(self);
+        let mem_start = self.mem_block.start_ptr().raw_addr(self);
+        let mem_end = self.mem_block.mapped_end_ptr().raw_addr(self);
         assert!(mem_start <= start_addr.raw_addr(self));
         assert!(start_addr.raw_addr(self) <= end_addr.raw_addr(self));
         assert!(end_addr.raw_addr(self) <= mem_end);
 
         // Ignore empty code ranges
         if start_addr == end_addr {
-            return (0..0).into_iter();
+            return 0..0;
         }
 
         let start_page = (start_addr.raw_addr(self) - mem_start) / self.page_size;
         let end_page = (end_addr.raw_addr(self) - mem_start - 1) / self.page_size;
-        (start_page..end_page + 1).into_iter()
+        start_page..end_page + 1
     }
 
     /// Get a (possibly dangling) direct pointer to the current write position
@@ -458,7 +457,7 @@ impl CodeBlock {
     /// Write a single byte at the current position.
     pub fn write_byte(&mut self, byte: u8) {
         let write_ptr = self.get_write_ptr();
-        if self.has_capacity(1) && self.mem_block.borrow_mut().write_byte(write_ptr, byte).is_ok() {
+        if self.has_capacity(1) && self.mem_block.write_byte(write_ptr, byte).is_ok() {
             self.write_pos += 1;
         } else {
             self.dropped_bytes = true;
@@ -590,8 +589,12 @@ impl CodeBlock {
         self.label_refs = state.label_refs;
     }
 
+    pub fn mark_all_writeable(&mut self) {
+        self.mem_block.mark_all_writeable();
+    }
+
     pub fn mark_all_executable(&mut self) {
-        self.mem_block.borrow_mut().mark_all_executable();
+        self.mem_block.mark_all_executable();
     }
 
     /// Code GC. Free code pages that are not on stack and reuse them.
@@ -689,7 +692,7 @@ impl CodeBlock {
         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);
 
-        Self::new(Rc::new(RefCell::new(virt_mem)), false, Rc::new(None), true)
+        Self::new(Rc::new(virt_mem), false, Rc::new(None), true)
     }
 
     /// Stubbed CodeBlock for testing conditions that can arise due to code GC. Can't execute generated code.
@@ -707,7 +710,7 @@ impl CodeBlock {
         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);
 
-        Self::new(Rc::new(RefCell::new(virt_mem)), false, Rc::new(Some(freed_pages)), true)
+        Self::new(Rc::new(virt_mem), false, Rc::new(Some(freed_pages)), true)
     }
 }
 
@@ -715,7 +718,7 @@ impl CodeBlock {
 impl fmt::LowerHex for CodeBlock {
     fn fmt(&self, fmtr: &mut fmt::Formatter) -> fmt::Result {
         for pos in 0..self.write_pos {
-            let mem_block = &*self.mem_block.borrow();
+            let mem_block = &*self.mem_block;
             let byte = unsafe { mem_block.start_ptr().raw_ptr(mem_block).add(pos).read() };
             fmtr.write_fmt(format_args!("{:02x}", byte))?;
         }
@@ -725,7 +728,7 @@ impl fmt::LowerHex for CodeBlock {
 
 impl crate::virtualmem::CodePtrBase for CodeBlock {
     fn base_ptr(&self) -> std::ptr::NonNull<u8> {
-        self.mem_block.borrow().base_ptr()
+        self.mem_block.base_ptr()
     }
 }
 
diff --git a/yjit/src/asm/x86_64/mod.rs b/yjit/src/asm/x86_64/mod.rs
index fbbfa714d8..0ef5e92117 100644
--- a/yjit/src/asm/x86_64/mod.rs
+++ b/yjit/src/asm/x86_64/mod.rs
@@ -1027,7 +1027,10 @@ pub fn mov(cb: &mut CodeBlock, dst: X86Opnd, src: X86Opnd) {
             }
 
             let output_num_bits:u32 = if mem.num_bits > 32 { 32 } else { mem.num_bits.into() };
-            assert!(imm_num_bits(imm.value) <= (output_num_bits as u8));
+            assert!(
+                mem.num_bits < 64 || imm_num_bits(imm.value) <= (output_num_bits as u8),
+                "immediate value should be small enough to survive sign extension"
+            );
             cb.write_int(imm.value as u64, output_num_bits);
         },
         // M + UImm
@@ -1042,7 +1045,10 @@ pub fn mov(cb: &mut CodeBlock, dst: X86Opnd, src: X86Opnd) {
             }
 
             let output_num_bits = if mem.num_bits > 32 { 32 } else { mem.num_bits.into() };
-            assert!(imm_num_bits(uimm.value as i64) <= (output_num_bits as u8));
+            assert!(
+                mem.num_bits < 64 || imm_num_bits(uimm.value as i64) <= (output_num_bits as u8),
+                "immediate value should be small enough to survive sign extension"
+            );
             cb.write_int(uimm.value, output_num_bits);
         },
         // * + Imm/UImm
diff --git a/yjit/src/asm/x86_64/tests.rs b/yjit/src/asm/x86_64/tests.rs
index 5ae983270f..eefcbfd52e 100644
--- a/yjit/src/asm/x86_64/tests.rs
+++ b/yjit/src/asm/x86_64/tests.rs
@@ -193,6 +193,7 @@ fn test_mov() {
     check_bytes("48c7470801000000", |cb| mov(cb, mem_opnd(64, RDI, 8), imm_opnd(1)));
     //check_bytes("67c7400411000000", |cb| mov(cb, mem_opnd(32, EAX, 4), imm_opnd(0x34))); // We don't distinguish between EAX and RAX here - that's probably fine?
     check_bytes("c7400411000000", |cb| mov(cb, mem_opnd(32, RAX, 4), imm_opnd(17)));
+    check_bytes("c7400401000080", |cb| mov(cb, mem_opnd(32, RAX, 4), uimm_opnd(0x80000001)));
     check_bytes("41895814", |cb| mov(cb, mem_opnd(32, R8, 20), EBX));
     check_bytes("4d8913", |cb| mov(cb, mem_opnd(64, R11, 0), R10));
     check_bytes("48c742f8f4ffffff", |cb| mov(cb, mem_opnd(64, RDX, -8), imm_opnd(-12)));