YJIT: Port next_page to x86_64

k0kubun · k0kubun · commit 29f17dd3e4d6 · 2022-10-06T17:18:24.000-07:00
diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
@@ -703,8 +703,7 @@ impl Assembler
         // For each instruction
         let start_write_pos = cb.get_write_pos();
         let mut i: usize = 0;
-        while i < self.insns.len() {
-            let insn = &self.insns[i];
+        while let Some(insn) = self.insns.get(i) {
             let had_dropped_bytes = cb.has_dropped_bytes();
             let src_ptr = cb.get_write_ptr();
             match insn {
@@ -1003,7 +1002,7 @@ impl Assembler
                 }
             };
 
-            // On failure, jump to the next page and retry the last insn
+            // On failure, jump to the next page and retry the current insn
             if !had_dropped_bytes && cb.has_dropped_bytes() && cb.next_page(None) {
                 let dst_ptr = cb.get_write_ptr();
                 cb.set_write_ptr(src_ptr);
diff --git a/yjit/src/backend/tests.rs b/yjit/src/backend/tests.rs
@@ -253,7 +253,7 @@ fn test_jcc_ptr()
 {
     let (mut asm, mut cb) = setup_asm();
 
-    let side_exit = Target::CodePtr((5 as *mut u8).into());
+    let side_exit = Target::CodePtr(((cb.get_write_ptr().raw_ptr() as usize + 5) as *mut u8).into());
     let not_mask = asm.not(Opnd::mem(32, EC, RUBY_OFFSET_EC_INTERRUPT_MASK));
     asm.test(
         Opnd::mem(32, EC, RUBY_OFFSET_EC_INTERRUPT_FLAG),
@@ -270,7 +270,7 @@ fn test_jmp_ptr()
 {
     let (mut asm, mut cb) = setup_asm();
 
-    let stub = Target::CodePtr((5 as *mut u8).into());
+    let stub = Target::CodePtr(((cb.get_write_ptr().raw_ptr() as usize + 5) as *mut u8).into());
     asm.jmp(stub);
 
     asm.compile_with_num_regs(&mut cb, 0);
@@ -281,7 +281,7 @@ fn test_jo()
 {
     let (mut asm, mut cb) = setup_asm();
 
-    let side_exit = Target::CodePtr((5 as *mut u8).into());
+    let side_exit = Target::CodePtr(((cb.get_write_ptr().raw_ptr() as usize + 5) as *mut u8).into());
 
     let arg1 = Opnd::mem(64, SP, 0);
     let arg0 = Opnd::mem(64, SP, 8);
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
@@ -372,7 +372,10 @@ impl Assembler
 
         // For each instruction
         let start_write_pos = cb.get_write_pos();
-        for insn in &self.insns {
+        let mut i: usize = 0;
+        while let Some(insn) = self.insns.get(i) {
+            let had_dropped_bytes = cb.has_dropped_bytes();
+            let src_ptr = cb.get_write_ptr();
             match insn {
                 Insn::Comment(text) => {
                     if cfg!(feature = "asm_comments") {
@@ -662,6 +665,18 @@ impl Assembler
                 // instructions can never make it to the emit stage.
                 _ => panic!("unsupported instruction passed to x86 backend: {:?}", insn)
             };
+
+            // On failure, jump to the next page and retry the current insn
+            if !had_dropped_bytes && cb.has_dropped_bytes() && cb.next_page(None) {
+                let dst_ptr = cb.get_write_ptr();
+                cb.set_write_ptr(src_ptr);
+                cb.without_page_end_reserve(|cb| {
+                    jmp_ptr(cb, dst_ptr);
+                });
+                cb.set_write_ptr(dst_ptr);
+            } else {
+                i += 1;
+            }
         }
 
         gc_offsets
