Move edits out of Env (#143)

* Move edits out of Env

Return move edits from `resolve_inserted_moves` instead of eagerly
allocating that vector in `Env`.

* Move the final edit iteration into a method on Edits

* Make fields on Edits private, and add a sort method

Author: elliottt

src/ion/data_structures.rs

@@ -433,7 +433,6 @@ pub struct Env<'a, F: Function> {
     pub multi_fixed_reg_fixups: Vec<MultiFixedRegFixup>,
 
     // Output:
-    pub edits: Vec<(PosWithPrio, Edit)>,
     pub allocs: Vec<Allocation>,
     pub inst_alloc_offsets: Vec<u32>,
     pub num_spillslots: u32,
@@ -692,6 +691,51 @@ impl InsertedMoves {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct Edits {
+    edits: Vec<(PosWithPrio, Edit)>,
+}
+
+impl Edits {
+    #[inline(always)]
+    pub fn with_capacity(n: usize) -> Self {
+        Self {
+            edits: Vec::with_capacity(n),
+        }
+    }
+
+    #[inline(always)]
+    pub fn len(&self) -> usize {
+        self.edits.len()
+    }
+
+    #[inline(always)]
+    pub fn iter(&self) -> impl Iterator<Item = &(PosWithPrio, Edit)> {
+        self.edits.iter()
+    }
+
+    #[inline(always)]
+    pub fn into_edits(self) -> impl Iterator<Item = (ProgPoint, Edit)> {
+        self.edits.into_iter().map(|(pos, edit)| (pos.pos, edit))
+    }
+
+    /// Sort edits by the combination of their program position and priority. This is a stable sort
+    /// to preserve the order of the moves the parallel move resolver inserts.
+    #[inline(always)]
+    pub fn sort(&mut self) {
+        self.edits.sort_by_key(|&(pos_prio, _)| pos_prio.key());
+    }
+
+    pub fn add(&mut self, pos_prio: PosWithPrio, from: Allocation, to: Allocation) {
+        if from != to {
+            if from.is_reg() && to.is_reg() {
+                debug_assert_eq!(from.as_reg().unwrap().class(), to.as_reg().unwrap().class());
+            }
+            self.edits.push((pos_prio, Edit::Move { from, to }));
+        }
+    }
+}
+
 /// The fields in this struct are reversed in sort order so that the entire
 /// struct can be treated as a u64 for sorting purposes.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]

src/ion/mod.rs

@@ -74,7 +74,6 @@ impl<'a, F: Function> Env<'a, F> {
             preferred_victim_by_class: [PReg::invalid(), PReg::invalid(), PReg::invalid()],
 
             multi_fixed_reg_fixups: vec![],
-            edits: Vec::with_capacity(n),
             allocs: Vec::with_capacity(4 * n),
             inst_alloc_offsets: vec![],
             num_spillslots: 0,
@@ -103,14 +102,14 @@ impl<'a, F: Function> Env<'a, F> {
         Ok(())
     }
 
-    pub(crate) fn run(&mut self) -> Result<(), RegAllocError> {
+    pub(crate) fn run(&mut self) -> Result<Edits, RegAllocError> {
         self.process_bundles()?;
         self.try_allocating_regs_for_spilled_bundles();
         self.allocate_spillslots();
         let moves = self.apply_allocations_and_insert_moves();
-        self.resolve_inserted_moves(moves);
+        let edits = self.resolve_inserted_moves(moves);
         self.compute_stackmaps();
-        Ok(())
+        Ok(edits)
     }
 }
 
@@ -129,18 +128,14 @@ pub fn run<F: Function>(
     let mut env = Env::new(func, mach_env, cfginfo, enable_annotations);
     env.init()?;
 
-    env.run()?;
+    let edits = env.run()?;
 
     if enable_annotations {
         env.dump_results();
     }
 
     Ok(Output {
-        edits: env
-            .edits
-            .into_iter()
-            .map(|(pos_prio, edit)| (pos_prio.pos, edit))
-            .collect(),
+        edits: edits.into_edits().collect(),
         allocs: env.allocs,
         inst_alloc_offsets: env.inst_alloc_offsets,
         num_spillslots: env.num_spillslots as usize,

src/ion/moves.rs

@@ -17,8 +17,8 @@ use super::{
     RedundantMoveEliminator, VRegIndex, SLOT_NONE,
 };
 use crate::ion::data_structures::{
-    u64_key, BlockparamIn, BlockparamOut, CodeRange, FixedRegFixupLevel, LiveRangeKey,
-    LiveRangeListEntry, PosWithPrio,
+    u64_key, BlockparamIn, BlockparamOut, CodeRange, Edits, FixedRegFixupLevel, LiveRangeKey,
+    LiveRangeListEntry,
 };
 use crate::ion::reg_traversal::RegTraversalIter;
 use crate::moves::{MoveAndScratchResolver, ParallelMoves};
@@ -772,7 +772,7 @@ impl<'a, F: Function> Env<'a, F> {
         inserted_moves
     }
 
-    pub fn resolve_inserted_moves(&mut self, mut inserted_moves: InsertedMoves) {
+    pub fn resolve_inserted_moves(&mut self, mut inserted_moves: InsertedMoves) -> Edits {
         // For each program point, gather all moves together. Then
         // resolve (see cases below).
         let mut i = 0;
@@ -839,6 +839,7 @@ impl<'a, F: Function> Env<'a, F> {
         }
 
         let mut last_pos = ProgPoint::before(Inst::new(0));
+        let mut edits = Edits::with_capacity(self.func.num_insts());
 
         while i < inserted_moves.moves.len() {
             let start = i;
@@ -982,7 +983,7 @@ impl<'a, F: Function> Env<'a, F> {
                     trace!("  resolved: {} -> {} ({:?})", src, dst, to_vreg);
                     let action = redundant_moves.process_move(src, dst, to_vreg);
                     if !action.elide {
-                        self.add_move_edit(pos_prio, src, dst);
+                        edits.add(pos_prio, src, dst);
                     } else {
                         trace!("    -> redundant move elided");
                     }
@@ -994,28 +995,20 @@ impl<'a, F: Function> Env<'a, F> {
         // be a stable sort! We have to keep the order produced by the
         // parallel-move resolver for all moves within a single sort
         // key.
-        self.edits.sort_by_key(|&(pos_prio, _)| pos_prio.key());
-        self.stats.edits_count = self.edits.len();
+        edits.sort();
+        self.stats.edits_count = edits.len();
 
         // Add debug annotations.
         if self.annotations_enabled {
-            for i in 0..self.edits.len() {
-                let &(pos_prio, ref edit) = &self.edits[i];
+            for &(pos_prio, ref edit) in edits.iter() {
                 match edit {
                     &Edit::Move { from, to } => {
                         self.annotate(pos_prio.pos, format!("move {} -> {}", from, to));
                     }
                 }
             }
         }
-    }
 
-    pub fn add_move_edit(&mut self, pos_prio: PosWithPrio, from: Allocation, to: Allocation) {
-        if from != to {
-            if from.is_reg() && to.is_reg() {
-                debug_assert_eq!(from.as_reg().unwrap().class(), to.as_reg().unwrap().class());
-            }
-            self.edits.push((pos_prio, Edit::Move { from, to }));
-        }
+        edits
     }
 }