From f1b188143b0255cef498ce4fb7a331daca64e063 Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Tue, 21 Jun 2022 17:01:26 -0400
Subject: [PATCH] Fix backend transform bug, add test

---
 yjit/src/backend/ir.rs    | 20 +++++++++++++-------
 yjit/src/backend/tests.rs | 11 +++++++++--
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index bacbbd541d..66a498fb30 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -324,7 +324,7 @@ impl Assembler
                 Opnd::InsnOut{ idx, .. } => {
                     self.live_ranges[*idx] = insn_idx;
                 }
-                Opnd::Mem( Mem { base: MemBase::InsnOut(idx), .. }) => {
+                Opnd::Mem(Mem { base: MemBase::InsnOut(idx), .. }) => {
                     self.live_ranges[*idx] = insn_idx;
                 }
                 _ => {}
@@ -424,17 +424,21 @@ impl Assembler
             label_names: self.label_names,
         };
 
-        // indices maps from the old instruction index to the new instruction
+        // Indices maps from the old instruction index to the new instruction
         // index.
         let mut indices: Vec<usize> = Vec::default();
 
         // Map an operand to the next set of instructions by correcting previous
         // InsnOut indices.
         fn map_opnd(opnd: Opnd, indices: &mut Vec<usize>) -> Opnd {
-            if let Opnd::InsnOut{ idx, num_bits } = opnd {
-                Opnd::InsnOut{ idx: indices[idx], num_bits }
-            } else {
-                opnd
+            match opnd {
+                Opnd::InsnOut{ idx, num_bits } => {
+                    Opnd::InsnOut{ idx: indices[idx], num_bits }
+                }
+                Opnd::Mem(Mem{ base: MemBase::InsnOut(idx), disp, num_bits,  }) => {
+                    Opnd::Mem(Mem{ base:MemBase::InsnOut(indices[idx]), disp, num_bits })
+                }
+                _ => opnd
             }
         }
 
@@ -531,6 +535,8 @@ impl Assembler
     /// instruction. This is our implementation of the linear scan algorithm.
     pub(super) fn alloc_regs(mut self, regs: Vec<Reg>) -> Assembler
     {
+        //dbg!(&self);
+
         // First, create the pool of registers.
         let mut pool: u32 = 0;
 
@@ -585,7 +591,7 @@ impl Assembler
                             if let Opnd::Reg(reg) = asm.insns[start_index].out {
                                 dealloc_reg(&mut pool, &regs, &reg);
                             } else {
-                                unreachable!("no register allocated for insn");
+                                unreachable!("no register allocated for insn {:?}", op);
                             }
                         }
                     }
diff --git a/yjit/src/backend/tests.rs b/yjit/src/backend/tests.rs
index 3a0f14e1f4..6545d01517 100644
--- a/yjit/src/backend/tests.rs
+++ b/yjit/src/backend/tests.rs
@@ -173,12 +173,19 @@ fn test_base_insn_out()
 {
     let (mut asm, mut cb) = setup_asm();
 
+    // Forced register to be reused
+    // This also causes the insn sequence to change length
+    asm.mov(
+        Opnd::mem(64, SP, 8),
+        Opnd::mem(64, SP, 0)
+    );
+
     // Load the pointer into a register
-    let ptr_reg = asm.load(Opnd::const_ptr(0 as *const u8));
+    let ptr_reg = asm.load(Opnd::const_ptr(4351776248 as *const u8));
     let counter_opnd = Opnd::mem(64, ptr_reg, 0);
 
     // Increment and store the updated value
-    asm.incr_counter(counter_opnd, 1.into() );
+    asm.incr_counter(counter_opnd, 1.into());
 
     asm.compile_with_num_regs(&mut cb, 1);
 }