ZJIT: x86: split: Query new assembler for live ranges

Author: XrXr

zjit/src/backend/x86_64/mod.rs

@@ -128,6 +128,17 @@ impl Assembler
         let mut iterator = self.insns.into_iter().enumerate().peekable();
         let mut asm = Assembler::new_with_label_names(take(&mut self.label_names), live_ranges.len());
 
+        // When we add news instructions to `asm`, we also create new VRegs not in the old
+        // assembler. All VRegs with index larger than the old `live_ranges` live in the
+        // new assembler since all old VRegs are inherited and keep their index.
+        fn vreg_outlives_insn(vreg_idx: usize, older_live_ranges: &Vec<LiveRange>, new_asm: &Assembler, insn_idx: usize) -> bool {
+            older_live_ranges
+                .get(vreg_idx)
+                .map_or_else(|| new_asm.live_ranges.get(vreg_idx), Some)
+                .map(|live_range| live_range.end() > insn_idx)
+                .expect("VReg in old or new assembler")
+        }
+
         while let Some((index, mut insn)) = iterator.next() {
             let is_load = matches!(insn, Insn::Load { .. } | Insn::LoadInto { .. });
             let mut opnd_iter = insn.opnd_iter_mut();
@@ -183,7 +194,7 @@ impl Assembler
                                 },
                                 // Instruction output whose live range spans beyond this instruction
                                 (Opnd::VReg { idx, .. }, _) => {
-                                    if live_ranges[idx].end() > index {
+                                    if vreg_outlives_insn(idx, &live_ranges, &asm, index) {
                                         *left = asm.load(*left);
                                     }
                                 },
@@ -248,7 +259,7 @@ impl Assembler
                     match opnd {
                         // Instruction output whose live range spans beyond this instruction
                         Opnd::VReg { idx, .. } => {
-                            if live_ranges[*idx].end() > index {
+                            if vreg_outlives_insn(*idx, &live_ranges, &asm, index) {
                                 *opnd = asm.load(*opnd);
                             }
                         },
@@ -272,7 +283,7 @@ impl Assembler
                         // If we have an instruction output whose live range
                         // spans beyond this instruction, we have to load it.
                         Opnd::VReg { idx, .. } => {
-                            if live_ranges[idx].end() > index {
+                            if vreg_outlives_insn(idx, &live_ranges, &asm, index) {
                                 *truthy = asm.load(*truthy);
                             }
                         },
@@ -307,7 +318,7 @@ impl Assembler
                         // If we have an instruction output whose live range
                         // spans beyond this instruction, we have to load it.
                         Opnd::VReg { idx, .. } => {
-                            if live_ranges[idx].end() > index {
+                            if vreg_outlives_insn(idx, &live_ranges, &asm, index) {
                                 *opnd = asm.load(*opnd);
                             }
                         },