Fix live_ranges idx calculation (https://github.com/Shopify/ruby/pull/353)

forward_pass adjusts the indexes of our opnds to reflect the new
instructions as they are generated in the forward pass. However, we were
using the old live_ranges array, for which the new indexes are
incorrect.

This caused us to previously generate an IR which contained unnecessary
trivial load instructions (ex. mov rax, rax), because it was looking at
the wrong lifespans. Presumably this could also cause bugs because the
lifespan of the incorrectly considered operand idx could be short.

We've added an assert which would have failed on the previous trivial
case (but not necessarily all cases).

Co-authored-by: Matthew Draper <matthew@trebex.net>

2 files changed, 16 insertions, 5 deletions

diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 41842c9704..066e9dd4ce 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -567,7 +567,7 @@ impl Assembler
 
     /// Transform input instructions, consumes the input assembler
     pub(super) fn forward_pass<F>(mut self, mut map_insn: F) -> Assembler
-        where F: FnMut(&mut Assembler, usize, Op, Vec<Opnd>, Option<Target>, Option<String>, Option<PosMarkerFn>)
+        where F: FnMut(&mut Assembler, usize, Op, Vec<Opnd>, Option<Target>, Option<String>, Option<PosMarkerFn>, Vec<Opnd>)
     {
         let mut asm = Assembler {
             insns: Vec::default(),
@@ -594,6 +594,7 @@ impl Assembler
         }
 
         for (index, insn) in self.insns.drain(..).enumerate() {
+            let original_opnds = insn.opnds.clone();
             let opnds: Vec<Opnd> = insn.opnds.into_iter().map(|opnd| map_opnd(opnd, &mut indices)).collect();
 
             // For each instruction, either handle it here or allow the map_insn
@@ -603,7 +604,7 @@ impl Assembler
                     asm.comment(insn.text.unwrap().as_str());
                 },
                 _ => {
-                    map_insn(&mut asm, index, insn.op, opnds, insn.target, insn.text, insn.pos_marker);
+                    map_insn(&mut asm, index, insn.op, opnds, insn.target, insn.text, insn.pos_marker, original_opnds);
                 }
             };
 
@@ -664,7 +665,7 @@ impl Assembler
 
         let live_ranges: Vec<usize> = std::mem::take(&mut self.live_ranges);
 
-        let asm = self.forward_pass(|asm, index, op, opnds, target, text, pos_marker| {
+        let asm = self.forward_pass(|asm, index, op, opnds, target, text, pos_marker, original_insns| {
             // Check if this is the last instruction that uses an operand that
             // spans more than one instruction. In that case, return the
             // allocated register to the pool.
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 7074c8980b..c3885be811 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -94,7 +94,7 @@ impl Assembler
     {
         let live_ranges: Vec<usize> = std::mem::take(&mut self.live_ranges);
 
-        self.forward_pass(|asm, index, op, opnds, target, text, pos_marker| {
+        self.forward_pass(|asm, index, op, opnds, target, text, pos_marker, original_opnds| {
             // Load heap object operands into registers because most
             // instructions can't directly work with 64-bit constants
             let opnds = match op {
@@ -134,7 +134,17 @@ impl Assembler
                             }
                         },
                         // Instruction output whose live range spans beyond this instruction
-                        (Opnd::InsnOut { idx, .. }, _) => {
+                        (Opnd::InsnOut { .. }, _) => {
+                            let idx = match original_opnds[0] {
+                                Opnd::InsnOut { idx, .. } => {
+                                    idx
+                                },
+                                _ => panic!("nooooo")
+                            };
+
+                            // Our input must be from a previous instruction!
+                            assert!(idx < index);
+
                             if live_ranges[idx] > index {
                                 (asm.load(opnds[0]), opnds[1])
                             } else {