YJIT: Simplify code page switching logic, remove an assert

We have received a report of `assert!( !cb.has_dropped_bytes())` in
set_page() failing. The only explanation for this seems to be memory
allocation failing in write_byte(). The if condition implies that
`current_write_pos < dst_pos < mem_size`, which rules out failing to
encode the relative jump. The has_capacity() assert above not tripping
implies that we were in a place in the page where write_byte() did
attempt to write the byte and potentially made a syscall in the process.

Remove the assert, since memory allocation could fail. Also, return
failure if the destination is outside of the code region to detect that
out-of-memory situation quicker.

1 files changed, 8 insertions, 9 deletions

diff --git a/yjit/src/asm/mod.rs b/yjit/src/asm/mod.rs
index a89be1a2f8..2bd776ec81 100644
--- a/yjit/src/asm/mod.rs
+++ b/yjit/src/asm/mod.rs
@@ -189,22 +189,21 @@ impl CodeBlock {
         // but you need to waste some space for keeping write_pos for every single page.
         // It doesn't seem necessary for performance either. So we're currently not doing it.
         let dst_pos = self.get_page_pos(page_idx);
-        if self.page_size * page_idx < self.mem_size && self.write_pos < dst_pos {
+        if self.write_pos < dst_pos {
+            // Fail if next page is out of bounds
+            if dst_pos >= self.mem_size {
+                return false;
+            }
+
             // Reset dropped_bytes
             self.dropped_bytes = false;
 
-            // Convert dst_pos to dst_ptr
-            let src_pos = self.write_pos;
-            self.write_pos = dst_pos;
-            let dst_ptr = self.get_write_ptr();
-            self.write_pos = src_pos;
-            self.without_page_end_reserve(|cb| assert!(cb.has_capacity(cb.jmp_ptr_bytes())));
-
             // Generate jmp_ptr from src_pos to dst_pos
+            let dst_ptr = self.get_ptr(dst_pos);
             self.without_page_end_reserve(|cb| {
+                assert!(cb.has_capacity(cb.jmp_ptr_bytes()));
                 cb.add_comment("jump to next page");
                 jmp_ptr(cb, dst_ptr);
-                assert!(!cb.has_dropped_bytes());
             });
 
             // Update past_page_bytes for code_size() if this is a new page