Iterator (https://github.com/Shopify/ruby/pull/372)

* Iterator

* Use the new iterator for the X86 backend split

* Use iterator for reg alloc, remove forward pass

* Fix up iterator usage on AArch64

* Update yjit/src/backend/ir.rs

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

* Various PR feedback for iterators for IR

* Use a local mutable reference for a64_split

* Move tests from ir.rs to tests.rs in backend

* Fix x86 shift instructions live range calculation

* Iterator

* Use the new iterator for the X86 backend split

* Fix up x86 iterator usage

* Fix ARM iterator usage

* Remove unintentionally duplicated tests

4 files changed, 274 insertions, 150 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 234339ca4e..fac77f972d 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -182,12 +182,14 @@ impl Assembler
             }
         }
 
-        self.forward_pass(|asm, index, op, opnds, target, text, pos_marker, original_opnds| {
-            // Load all Value operands into registers that aren't already a part
-            // of Load instructions.
-            let opnds = match op {
-                Op::Load => opnds,
-                _ => opnds.into_iter().map(|opnd| {
+        let mut asm_local = Assembler::new_with_label_names(std::mem::take(&mut self.label_names));
+        let asm = &mut asm_local;
+        let mut iterator = self.into_draining_iter();
+
+        while let Some((index, insn)) = iterator.next_mapped() {
+            let opnds = match insn.op {
+                Op::Load => insn.opnds,
+                _ => insn.opnds.into_iter().map(|opnd| {
                     if let Opnd::Value(_) = opnd {
                         asm.load(opnd)
                     } else {
@@ -196,7 +198,7 @@ impl Assembler
                 }).collect()
             };
 
-            match op {
+            match insn.op {
                 Op::Add => {
                     match (opnds[0], opnds[1]) {
                         (Opnd::Reg(_) | Opnd::InsnOut { .. }, Opnd::Reg(_) | Opnd::InsnOut { .. }) => {
@@ -217,17 +219,17 @@ impl Assembler
                 Op::And | Op::Or => {
                     match (opnds[0], opnds[1]) {
                         (Opnd::Reg(_), Opnd::Reg(_)) => {
-                            asm.push_insn(op, vec![opnds[0], opnds[1]], target, text, pos_marker);
+                            asm.push_insn(insn.op, vec![opnds[0], opnds[1]], insn.target, insn.text, insn.pos_marker);
                         },
                         (reg_opnd @ Opnd::Reg(_), other_opnd) |
                         (other_opnd, reg_opnd @ Opnd::Reg(_)) => {
                             let opnd1 = split_bitmask_immediate(asm, other_opnd);
-                            asm.push_insn(op, vec![reg_opnd, opnd1], target, text, pos_marker);
+                            asm.push_insn(insn.op, vec![reg_opnd, opnd1], insn.target, insn.text, insn.pos_marker);
                         },
                         _ => {
                             let opnd0 = split_load_operand(asm, opnds[0]);
                             let opnd1 = split_bitmask_immediate(asm, opnds[1]);
-                            asm.push_insn(op, vec![opnd0, opnd1], target, text, pos_marker);
+                            asm.push_insn(insn.op, vec![opnd0, opnd1], insn.target, insn.text, insn.pos_marker);
                         }
                     }
                 },
@@ -246,7 +248,7 @@ impl Assembler
 
                     // Now we push the CCall without any arguments so that it
                     // just performs the call.
-                    asm.ccall(target.unwrap().unwrap_fun_ptr(), vec![]);
+                    asm.ccall(insn.target.unwrap().unwrap_fun_ptr(), vec![]);
                 },
                 Op::Cmp => {
                     let opnd0 = match opnds[0] {
@@ -273,7 +275,7 @@ impl Assembler
                         }
                     }).collect();
 
-                    asm.push_insn(op, new_opnds, target, text, pos_marker);
+                    asm.push_insn(insn.op, new_opnds, insn.target, insn.text, insn.pos_marker);
                 },
                 Op::IncrCounter => {
                     // We'll use LDADD later which only works with registers
@@ -392,10 +394,14 @@ impl Assembler
                     asm.test(opnd0, opnd1);
                 },
                 _ => {
-                    asm.push_insn(op, opnds, target, text, pos_marker);
+                    asm.push_insn(insn.op, opnds, insn.target, insn.text, insn.pos_marker);
                 }
             };
-        })
+
+            iterator.map_insn_index(asm);
+        }
+
+        asm_local
     }
 
     /// Emit platform-specific machine code
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 5eee61b228..2dfb859fe9 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -2,8 +2,10 @@
 #![allow(unused_variables)]
 #![allow(unused_imports)]
 
+use std::cell::Cell;
 use std::fmt;
 use std::convert::From;
+use std::mem::take;
 use crate::cruby::{VALUE};
 use crate::virtualmem::{CodePtr};
 use crate::asm::{CodeBlock, uimm_num_bits, imm_num_bits};
@@ -288,6 +290,20 @@ impl Opnd
             _ => unreachable!()
         }
     }
+
+    /// Maps the indices from a previous list of instructions to a new list of
+    /// instructions.
+    pub fn map_index(self, indices: &Vec<usize>) -> Opnd {
+        match self {
+            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 })
+            },
+            _ => self
+        }
+    }
 }
 
 impl From<usize> for Opnd {
@@ -433,11 +449,15 @@ pub struct Assembler
 
 impl Assembler
 {
-    pub fn new() -> Assembler {
-        Assembler {
+    pub fn new() -> Self {
+        Self::new_with_label_names(Vec::default())
+    }
+
+    pub fn new_with_label_names(label_names: Vec<String>) -> Self {
+        Self {
             insns: Vec::default(),
             live_ranges: Vec::default(),
-            label_names: Vec::default(),
+            label_names
         }
     }
 
@@ -573,58 +593,6 @@ impl Assembler
         self.live_ranges.push(self.insns.len());
     }
 
-    /// 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>, Vec<Opnd>)
-    {
-        let mut asm = Assembler {
-            insns: Vec::default(),
-            live_ranges: Vec::default(),
-            label_names: self.label_names,
-        };
-
-        // 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 {
-            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
-            }
-        }
-
-        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
-            // callback to handle it.
-            match insn.op {
-                Op::Comment => {
-                    asm.comment(insn.text.unwrap().as_str());
-                },
-                _ => {
-                    map_insn(&mut asm, index, insn.op, opnds, insn.target, insn.text, insn.pos_marker, original_opnds);
-                }
-            };
-
-            // Here we're assuming that if we've pushed multiple instructions,
-            // the output that we're using is still the final instruction that
-            // was pushed.
-            indices.push(asm.insns.len() - 1);
-        }
-
-        asm
-    }
-
     /// Sets the out field on the various instructions that require allocated
     /// registers because their output is used as the operand on a subsequent
     /// instruction. This is our implementation of the linear scan algorithm.
@@ -671,13 +639,15 @@ impl Assembler
             }
         }
 
-        let live_ranges: Vec<usize> = std::mem::take(&mut self.live_ranges);
+        let live_ranges: Vec<usize> = take(&mut self.live_ranges);
+        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
+        let mut iterator = self.into_draining_iter();
 
-        let asm = self.forward_pass(|asm, index, op, opnds, target, text, pos_marker, original_insns| {
+        while let Some((index, insn)) = iterator.next_unmapped() {
             // 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.
-            for opnd in &opnds {
+            for opnd in &insn.opnds {
                 match opnd {
                     Opnd::InsnOut{idx, .. } |
                     Opnd::Mem( Mem { base: MemBase::InsnOut(idx), .. }) => {
@@ -693,7 +663,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 {:?}", op);
+                                unreachable!("no register allocated for insn {:?}", insn.op);
                             }
                         }
                     }
@@ -703,7 +673,7 @@ impl Assembler
             }
 
             // C return values need to be mapped to the C return register
-            if op == Op::CCall {
+            if insn.op == Op::CCall {
                 assert_eq!(pool, 0, "register lives past C function call");
             }
 
@@ -713,7 +683,7 @@ impl Assembler
             if live_ranges[index] != index {
 
                 // C return values need to be mapped to the C return register
-                if op == Op::CCall {
+                if insn.op == Op::CCall {
                     out_reg = Opnd::Reg(take_reg(&mut pool, &regs, &C_RET_REG))
                 }
 
@@ -722,8 +692,8 @@ impl Assembler
                 // We do this to improve register allocation on x86
                 // e.g. out  = add(reg0, reg1)
                 //      reg0 = add(reg0, reg1)
-                else if opnds.len() > 0 {
-                    if let Opnd::InsnOut{idx, ..} = opnds[0] {
+                else if insn.opnds.len() > 0 {
+                    if let Opnd::InsnOut{idx, ..} = insn.opnds[0] {
                         if live_ranges[idx] == index {
                             if let Opnd::Reg(reg) = asm.insns[idx].out {
                                 out_reg = Opnd::Reg(take_reg(&mut pool, &regs, &reg))
@@ -734,9 +704,9 @@ impl Assembler
 
                 // Allocate a new register for this instruction
                 if out_reg == Opnd::None {
-                    out_reg = if op == Op::LiveReg {
+                    out_reg = if insn.op == Op::LiveReg {
                         // Allocate a specific register
-                        let reg = opnds[0].unwrap_reg();
+                        let reg = insn.opnds[0].unwrap_reg();
                         Opnd::Reg(take_reg(&mut pool, &regs, &reg))
                     } else {
                         Opnd::Reg(alloc_reg(&mut pool, &regs))
@@ -745,7 +715,7 @@ impl Assembler
             }
 
             // Replace InsnOut operands by their corresponding register
-            let reg_opnds: Vec<Opnd> = opnds.into_iter().map(|opnd|
+            let reg_opnds: Vec<Opnd> = insn.opnds.into_iter().map(|opnd|
                 match opnd {
                     Opnd::InsnOut{idx, ..} => asm.insns[idx].out,
                     Opnd::Mem(Mem { base: MemBase::InsnOut(idx), disp, num_bits }) => {
@@ -760,7 +730,7 @@ impl Assembler
                 }
             ).collect();
 
-            asm.push_insn(op, reg_opnds, target, text, pos_marker);
+            asm.push_insn(insn.op, reg_opnds, insn.target, insn.text, insn.pos_marker);
 
             // Set the output register for this instruction
             let num_insns = asm.insns.len();
@@ -770,7 +740,7 @@ impl Assembler
                 out_reg = Opnd::Reg(reg.sub_reg(num_out_bits))
             }
             new_insn.out = out_reg;
-        });
+        }
 
         assert_eq!(pool, 0, "Expected all registers to be returned to the pool");
         asm
@@ -792,6 +762,123 @@ impl Assembler
         let alloc_regs = alloc_regs.drain(0..num_regs).collect();
         self.compile_with_regs(cb, alloc_regs)
     }
+
+    /// Consume the assembler by creating a new draining iterator.
+    pub fn into_draining_iter(self) -> AssemblerDrainingIterator {
+        AssemblerDrainingIterator::new(self)
+    }
+
+    /// Consume the assembler by creating a new lookback iterator.
+    pub fn into_lookback_iter(self) -> AssemblerLookbackIterator {
+        AssemblerLookbackIterator::new(self)
+    }
+
+    pub fn ccall(&mut self, fptr: *const u8, opnds: Vec<Opnd>) -> Opnd {
+        let target = Target::FunPtr(fptr);
+        self.push_insn(Op::CCall, opnds, Some(target), None, None)
+    }
+
+    // pub fn pos_marker<F: FnMut(CodePtr)>(&mut self, marker_fn: F)
+    pub fn pos_marker(&mut self, marker_fn: impl Fn(CodePtr) + 'static) {
+        self.push_insn(Op::PosMarker, vec![], None, None, Some(Box::new(marker_fn)));
+    }
+}
+
+/// A struct that allows iterating through an assembler's instructions and
+/// consuming them as it iterates.
+pub struct AssemblerDrainingIterator {
+    insns: std::vec::IntoIter<Insn>,
+    index: usize,
+    indices: Vec<usize>
+}
+
+impl AssemblerDrainingIterator {
+    fn new(asm: Assembler) -> Self {
+        Self {
+            insns: asm.insns.into_iter(),
+            index: 0,
+            indices: Vec::default()
+        }
+    }
+
+    /// When you're working with two lists of instructions, you need to make
+    /// sure you do some bookkeeping to align the indices contained within the
+    /// operands of the two lists.
+    ///
+    /// This function accepts the assembler that is being built and tracks the
+    /// end of the current list of instructions in order to maintain that
+    /// alignment.
+    pub fn map_insn_index(&mut self, asm: &mut Assembler) {
+        self.indices.push(asm.insns.len() - 1);
+    }
+
+    /// Map an operand by using this iterator's list of mapped indices.
+    pub fn map_opnd(&self, opnd: Opnd) -> Opnd {
+        opnd.map_index(&self.indices)
+    }
+
+    /// Returns the next instruction in the list with the indices corresponding
+    /// to the next list of instructions.
+    pub fn next_mapped(&mut self) -> Option<(usize, Insn)> {
+        self.next_unmapped().map(|(index, insn)| {
+            let opnds = insn.opnds.into_iter().map(|opnd| opnd.map_index(&self.indices)).collect();
+            (index, Insn { opnds, ..insn })
+        })
+    }
+
+    /// Returns the next instruction in the list with the indices corresponding
+    /// to the previous list of instructions.
+    pub fn next_unmapped(&mut self) -> Option<(usize, Insn)> {
+        let index = self.index;
+        self.index += 1;
+        self.insns.next().map(|insn| (index, insn))
+    }
+}
+
+/// A struct that allows iterating through references to an assembler's
+/// instructions without consuming them.
+pub struct AssemblerLookbackIterator {
+    asm: Assembler,
+    index: Cell<usize>
+}
+
+impl AssemblerLookbackIterator {
+    fn new(asm: Assembler) -> Self {
+        Self { asm, index: Cell::new(0) }
+    }
+
+    /// Fetches a reference to an instruction at a specific index.
+    pub fn get(&self, index: usize) -> Option<&Insn> {
+        self.asm.insns.get(index)
+    }
+
+    /// Fetches a reference to an instruction in the list relative to the
+    /// current cursor location of this iterator.
+    pub fn get_relative(&self, difference: i32) -> Option<&Insn> {
+        let index: Result<i32, _> = self.index.get().try_into();
+        let relative: Result<usize, _> = index.and_then(|value| (value + difference).try_into());
+        relative.ok().and_then(|value| self.asm.insns.get(value))
+    }
+
+    /// Fetches the previous instruction relative to the current cursor location
+    /// of this iterator.
+    pub fn get_previous(&self) -> Option<&Insn> {
+        self.get_relative(-1)
+    }
+
+    /// Fetches the next instruction relative to the current cursor location of
+    /// this iterator.
+    pub fn get_next(&self) -> Option<&Insn> {
+        self.get_relative(1)
+    }
+
+    /// Returns the next instruction in the list with the indices corresponding
+    /// to the previous list of instructions.
+    pub fn next_unmapped(&self) -> Option<(usize, &Insn)> {
+        let index = self.index.get();
+        self.index.set(index + 1);
+        self.asm.insns.get(index).map(|insn| (index, insn))
+    }
 }
 
 impl fmt::Debug for Assembler {
@@ -806,21 +893,6 @@ impl fmt::Debug for Assembler {
     }
 }
 
-impl Assembler
-{
-    pub fn ccall(&mut self, fptr: *const u8, opnds: Vec<Opnd>) -> Opnd
-    {
-        let target = Target::FunPtr(fptr);
-        self.push_insn(Op::CCall, opnds, Some(target), None, None)
-    }
-
-    //pub fn pos_marker<F: FnMut(CodePtr)>(&mut self, marker_fn: F)
-    pub fn pos_marker(&mut self, marker_fn: impl Fn(CodePtr) + 'static)
-    {
-        self.push_insn(Op::PosMarker, vec![], None, None, Some(Box::new(marker_fn)));
-    }
-}
-
 macro_rules! def_push_jcc {
     ($op_name:ident, $opcode:expr) => {
         impl Assembler
diff --git a/yjit/src/backend/tests.rs b/yjit/src/backend/tests.rs
index a31e16071b..e4ab95d4ee 100644
--- a/yjit/src/backend/tests.rs
+++ b/yjit/src/backend/tests.rs
@@ -299,3 +299,41 @@ fn test_bake_string() {
     asm.bake_string("Hello, world!");
     asm.compile_with_num_regs(&mut cb, 0);
 }
+
+#[test]
+fn test_draining_iterator() {
+    let mut asm = Assembler::new();
+
+    asm.load(Opnd::None);
+    asm.store(Opnd::None, Opnd::None);
+    asm.add(Opnd::None, Opnd::None);
+
+    let mut iter = asm.into_draining_iter();
+
+    while let Some((index, insn)) = iter.next_unmapped() {
+        match index {
+            0 => assert_eq!(insn.op, Op::Load),
+            1 => assert_eq!(insn.op, Op::Store),
+            2 => assert_eq!(insn.op, Op::Add),
+            _ => panic!("Unexpected instruction index"),
+        };
+    }
+}
+
+#[test]
+fn test_lookback_iterator() {
+    let mut asm = Assembler::new();
+
+    asm.load(Opnd::None);
+    asm.store(Opnd::None, Opnd::None);
+    asm.store(Opnd::None, Opnd::None);
+
+    let mut iter = asm.into_lookback_iter();
+
+    while let Some((index, insn)) = iter.next_unmapped() {
+        if index > 0 {
+            assert_eq!(iter.get_previous().unwrap().opnds[0], Opnd::None);
+            assert_eq!(insn.op, Op::Store);
+        }
+    }
+}
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index b0802b3187..9fcbb69a68 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -2,11 +2,13 @@
 #![allow(unused_variables)]
 #![allow(unused_imports)]
 
+use std::mem::take;
+
 use crate::asm::*;
 use crate::asm::x86_64::*;
 use crate::codegen::{JITState};
 use crate::cruby::*;
-use crate::backend::ir::{Assembler, Opnd, Target, Op, MemBase, Mem};
+use crate::backend::ir::*;
 
 // Use the x86 register type for this platform
 pub type Reg = X86Reg;
@@ -94,31 +96,51 @@ impl Assembler
     /// Split IR instructions for the x86 platform
     fn x86_split(mut self) -> Assembler
     {
-        let live_ranges: Vec<usize> = std::mem::take(&mut self.live_ranges);
-
-        self.forward_pass(|asm, index, op, opnds, target, text, pos_marker, original_opnds| {
-            // Load VALUEs into registers because
-            //  - Most instructions can't be encoded with 64-bit immediates.
-            //  - We look for Op::Load specifically when emiting to keep GC'ed
-            //    VALUEs alive. This is a sort of canonicalization.
-            let opnds = match op {
-                Op::Load => opnds,
-                _ => opnds.into_iter().map(|opnd| {
-                    if let Opnd::Value(value) = opnd {
-                        // Since mov(mem64, imm32) sign extends, as_i64() makes sure we split
-                        // when the extended value is different.
-                        if !value.special_const_p() || imm_num_bits(value.as_i64()) > 32 {
-                            return asm.load(opnd);
-                        }
+        let live_ranges: Vec<usize> = take(&mut self.live_ranges);
+        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
+        let mut iterator = self.into_draining_iter();
+
+        while let Some((index, insn)) = iterator.next_unmapped() {
+            // When we're iterating through the instructions with x86_split, we
+            // need to know the previous live ranges in order to tell if a
+            // register lasts beyond the current instruction. So instead of
+            // using next_mapped, we call next_unmapped. When you're using the
+            // next_unmapped API, you need to make sure that you map each
+            // operand that could reference an old index, which means both
+            // Opnd::InsnOut operands and Opnd::Mem operands with a base of
+            // MemBase::InsnOut.
+            //
+            // You need to ensure that you only map it _once_, because otherwise
+            // you'll end up mapping an incorrect index which could end up being
+            // out of bounds of the old set of indices.
+            //
+            // We handle all of that mapping here to ensure that it's only
+            // mapped once. We also handle loading Opnd::Value operands into
+            // registers here so that all mapping happens in one place. We load
+            // Opnd::Value operands into registers here because:
+            //
+            //   - Most instructions can't be encoded with 64-bit immediates.
+            //   - We look for Op::Load specifically when emiting to keep GC'ed
+            //     VALUEs alive. This is a sort of canonicalization.
+            let opnds: Vec<Opnd> = insn.opnds.iter().map(|opnd| {
+                if insn.op == Op::Load {
+                    iterator.map_opnd(*opnd)
+                } else if let Opnd::Value(value) = opnd {
+                    // Since mov(mem64, imm32) sign extends, as_i64() makes sure
+                    // we split when the extended value is different.
+                    if !value.special_const_p() || imm_num_bits(value.as_i64()) > 32 {
+                        asm.load(iterator.map_opnd(*opnd))
+                    } else {
+                        iterator.map_opnd(*opnd)
                     }
+                } else {
+                    iterator.map_opnd(*opnd)
+                }
+            }).collect();
 
-                    opnd
-                }).collect()
-            };
-
-            match op {
+            match insn.op {
                 Op::Add | Op::Sub | Op::And | Op::Cmp | Op::Or | Op::Test => {
-                    let (opnd0, opnd1) = match (opnds[0], opnds[1]) {
+                    let (opnd0, opnd1) = match (insn.opnds[0], insn.opnds[1]) {
                         (Opnd::Mem(_), Opnd::Mem(_)) => {
                             (asm.load(opnds[0]), asm.load(opnds[1]))
                         },
@@ -138,17 +160,7 @@ impl Assembler
                             }
                         },
                         // Instruction output whose live range spans beyond this instruction
-                        (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);
-
+                        (Opnd::InsnOut { idx, .. }, _) => {
                             if live_ranges[idx] > index {
                                 (asm.load(opnds[0]), opnds[1])
                             } else {
@@ -162,24 +174,14 @@ impl Assembler
                         _ => (opnds[0], opnds[1])
                     };
 
-                    asm.push_insn(op, vec![opnd0, opnd1], target, text, pos_marker);
+                    asm.push_insn(insn.op, vec![opnd0, opnd1], insn.target, insn.text, insn.pos_marker);
                 },
                 // These instructions modify their input operand in-place, so we
                 // may need to load the input value to preserve it
                 Op::LShift | Op::RShift | Op::URShift => {
-                    let (opnd0, opnd1) = match (opnds[0], opnds[1]) {
+                    let (opnd0, opnd1) = match (insn.opnds[0], insn.opnds[1]) {
                         // Instruction output whose live range spans beyond this instruction
-                        (Opnd::InsnOut { .. }, _) => {
-                            let idx = match original_opnds[0] {
-                                Opnd::InsnOut { idx, .. } => {
-                                    idx
-                                },
-                                _ => unreachable!()
-                            };
-
-                            // Our input must be from a previous instruction!
-                            assert!(idx < index);
-
+                        (Opnd::InsnOut { idx, .. }, _) => {
                             if live_ranges[idx] > index {
                                 (asm.load(opnds[0]), opnds[1])
                             } else {
@@ -193,7 +195,7 @@ impl Assembler
                         _ => (opnds[0], opnds[1])
                     };
 
-                    asm.push_insn(op, vec![opnd0, opnd1], target, text, pos_marker);
+                    asm.push_insn(insn.op, vec![opnd0, opnd1], insn.target, insn.text, insn.pos_marker);
                 },
                 Op::CSelZ | Op::CSelNZ | Op::CSelE | Op::CSelNE |
                 Op::CSelL | Op::CSelLE | Op::CSelG | Op::CSelGE => {
@@ -204,7 +206,7 @@ impl Assembler
                         }
                     }).collect();
 
-                    asm.push_insn(op, new_opnds, target, text, pos_marker);
+                    asm.push_insn(insn.op, new_opnds, insn.target, insn.text, insn.pos_marker);
                 },
                 Op::Mov => {
                     match (opnds[0], opnds[1]) {
@@ -236,7 +238,7 @@ impl Assembler
                     }
                 },
                 Op::Not => {
-                    let opnd0 = match opnds[0] {
+                    let opnd0 = match insn.opnds[0] {
                         // If we have an instruction output whose live range
                         // spans beyond this instruction, we have to load it.
                         Opnd::InsnOut { idx, .. } => {
@@ -248,7 +250,9 @@ impl Assembler
                         },
                         // We have to load memory and register operands to avoid
                         // corrupting them.
-                        Opnd::Mem(_) | Opnd::Reg(_) => asm.load(opnds[0]),
+                        Opnd::Mem(_) | Opnd::Reg(_) => {
+                            asm.load(opnds[0])
+                        },
                         // Otherwise we can just reuse the existing operand.
                         _ => opnds[0]
                     };
@@ -256,10 +260,14 @@ impl Assembler
                     asm.not(opnd0);
                 },
                 _ => {
-                    asm.push_insn(op, opnds, target, text, pos_marker);
+                    asm.push_insn(insn.op, opnds, insn.target, insn.text, insn.pos_marker);
                 }
             };
-        })
+
+            iterator.map_insn_index(&mut asm);
+        }
+
+        asm
     }
 
     /// Emit platform-specific machine code