Merge pull request #23 from Amanieu/iter

Add a helper to iterate over insts and edits of a block in order

src/checker.rs

@@ -77,8 +77,8 @@
 #![allow(dead_code)]
 
 use crate::{
-    Allocation, AllocationKind, Block, Edit, Function, Inst, InstPosition, Operand,
+    Allocation, AllocationKind, Block, Edit, Function, Inst, InstOrEdit, InstPosition, Operand,
-    OperandConstraint, OperandKind, OperandPos, Output, PReg, ProgPoint, RegClass, VReg,
+    OperandConstraint, OperandKind, OperandPos, Output, PReg, RegClass, VReg,
 };
 
 use std::collections::{HashMap, HashSet, VecDeque};
@@ -544,18 +544,30 @@ impl<'a, F: Function> Checker<'a, F> {
                 .push(slot);
         }
 
-        // For each original instruction, create an `Op`.
         let mut last_inst = None;
-        let mut insert_idx = 0;
         for block in 0..self.f.num_blocks() {
             let block = Block::new(block);
-            for inst in self.f.block_insns(block).iter() {
+            for inst_or_edit in out.block_insts_and_edits(self.f, block) {
+                match inst_or_edit {
+                    InstOrEdit::Inst(inst) => {
                         assert!(last_inst.is_none() || inst > last_inst.unwrap());
                         last_inst = Some(inst);
+                        self.handle_inst(block, inst, &mut safepoint_slots, out);
+                    }
+                    InstOrEdit::Edit(edit) => self.handle_edit(block, edit),
+                }
+            }
+        }
+    }
 
-                // Any inserted edits before instruction.
+    /// For each original instruction, create an `Op`.
-                self.handle_edits(block, out, &mut insert_idx, ProgPoint::before(inst));
+    fn handle_inst(
-
+        &mut self,
+        block: Block,
+        inst: Inst,
+        safepoint_slots: &mut HashMap<Inst, Vec<Allocation>>,
+        out: &Output,
+    ) {
         // If this is a safepoint, then check the spillslots at this point.
         if self.f.requires_refs_on_stack(inst) {
             let allocs = safepoint_slots.remove(&inst).unwrap_or_else(|| vec![]);
@@ -582,21 +594,10 @@ impl<'a, F: Function> Checker<'a, F> {
             log::trace!("checker: adding inst {:?}", checkinst);
             self.bb_insts.get_mut(&block).unwrap().push(checkinst);
         }
-
-                // Any inserted edits after instruction.
-                self.handle_edits(block, out, &mut insert_idx, ProgPoint::after(inst));
-            }
-        }
     }
 
-    fn handle_edits(&mut self, block: Block, out: &Output, idx: &mut usize, pos: ProgPoint) {
+    fn handle_edit(&mut self, block: Block, edit: &Edit) {
-        while *idx < out.edits.len() && out.edits[*idx].0 <= pos {
+        log::trace!("checker: adding edit {:?}", edit);
-            let &(edit_pos, ref edit) = &out.edits[*idx];
-            *idx += 1;
-            if edit_pos < pos {
-                continue;
-            }
-            log::trace!("checker: adding edit {:?} at pos {:?}", edit, pos);
         match edit {
             &Edit::Move { from, to, to_vreg } => {
                 self.bb_insts
@@ -618,7 +619,6 @@ impl<'a, F: Function> Checker<'a, F> {
             }
         }
     }
-    }
 
     /// Perform the dataflow analysis to compute checker state at each BB entry.
     fn analyze(&mut self) {

src/lib.rs

@@ -1122,6 +1122,46 @@ pub enum Edit {
     DefAlloc { alloc: Allocation, vreg: VReg },
 }
 
+/// Wrapper around either an original instruction or an inserted edit.
+#[derive(Clone, Debug)]
+pub enum InstOrEdit<'a> {
+    Inst(Inst),
+    Edit(&'a Edit),
+}
+
+/// Iterator over the instructions and edits in a block.
+pub struct OutputIter<'a> {
+    /// List of edits starting at the first for the current block.
+    edits: &'a [(ProgPoint, Edit)],
+
+    /// Remaining instructions in the current block.
+    inst_range: InstRange,
+}
+
+impl<'a> Iterator for OutputIter<'a> {
+    type Item = InstOrEdit<'a>;
+
+    fn next(&mut self) -> Option<InstOrEdit<'a>> {
+        // There can't be any edits after the last instruction in a block, so
+        // we don't need to worry about that case.
+        if self.inst_range.len() == 0 {
+            return None;
+        }
+
+        // Return any edits that happen before the next instruction first.
+        let next_inst = self.inst_range.first();
+        if let Some((edit, remaining_edits)) = self.edits.split_first() {
+            if edit.0 <= ProgPoint::before(next_inst) {
+                self.edits = remaining_edits;
+                return Some(InstOrEdit::Edit(&edit.1));
+            }
+        }
+
+        self.inst_range = self.inst_range.rest();
+        Some(InstOrEdit::Inst(next_inst))
+    }
+}
+
 /// A machine envrionment tells the register allocator which registers
 /// are available to allocate and what register may be used as a
 /// scratch register for each class, and some other miscellaneous info
@@ -1206,6 +1246,30 @@ impl Output {
         };
         &self.allocs[start..end]
     }
+
+    /// Returns an iterator over the instructions and edits in a block, in
+    /// order.
+    pub fn block_insts_and_edits(&self, func: &impl Function, block: Block) -> OutputIter<'_> {
+        let inst_range = func.block_insns(block);
+
+        let edit_idx = self
+            .edits
+            .binary_search_by(|&(pos, _)| {
+                // This predicate effectively searches for a point *just* before
+                // the first ProgPoint. This never returns Ordering::Equal, but
+                // binary_search_by returns the index of where it would have
+                // been inserted in Err.
+                if pos < ProgPoint::before(inst_range.first()) {
+                    std::cmp::Ordering::Less
+                } else {
+                    std::cmp::Ordering::Greater
+                }
+            })
+            .unwrap_err();
+
+        let edits = &self.edits[edit_idx..];
+        OutputIter { inst_range, edits }
+    }
 }
 
 /// An error that prevents allocation.