diff --git a/src/lib.rs b/src/lib.rs index 36b8042..9639b0f 100644 --- a/src/lib.rs +++ b/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> { + // 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.