Apply review feedback

Cargo.toml

@@ -11,6 +11,7 @@ repository = "https://github.com/bytecodealliance/regalloc2"
 log = { version = "0.4.8", default-features = false }
 smallvec = "1.6.1"
 fxhash = "0.2.1"
+slice-group-by = "0.3.0"
 
 # The below are only needed for fuzzing.
 # Keep this in sync with libfuzzer_sys's crate version:

src/ion/data_structures.rs

@@ -269,6 +269,15 @@ pub struct PRegData {
     pub is_stack: bool,
 }
 
+#[derive(Clone, Debug)]
+pub struct MultiFixedRegFixup {
+    pub pos: ProgPoint,
+    pub from_slot: u8,
+    pub to_slot: u8,
+    pub to_preg: PRegIndex,
+    pub vreg: VRegIndex,
+}
+
 #[derive(Clone, Debug)]
 pub struct Env<'a, F: Function> {
     pub func: &'a F,
@@ -330,9 +339,7 @@ pub struct Env<'a, F: Function> {
     // the register available. When we produce the final edit-list, we
     // will insert a copy from wherever the VReg's primary allocation
     // was to the approprate PReg.
-    //
+    pub multi_fixed_reg_fixups: Vec<MultiFixedRegFixup>,
-    // (progpoint, from-slot, copy-to-preg, vreg, to-slot)
-    pub multi_fixed_reg_fixups: Vec<(ProgPoint, u8, PRegIndex, VRegIndex, u8)>,
 
     pub inserted_moves: Vec<InsertedMove>,
 

src/ion/liveranges.rs

@@ -18,12 +18,13 @@ use super::{
     SpillSetIndex, Use, VRegData, VRegIndex, SLOT_NONE,
 };
 use crate::indexset::IndexSet;
-use crate::util::SliceGroupBy;
+use crate::ion::data_structures::MultiFixedRegFixup;
 use crate::{
     Allocation, Block, Function, Inst, InstPosition, Operand, OperandConstraint, OperandKind,
     OperandPos, PReg, ProgPoint, RegAllocError, VReg,
 };
 use fxhash::FxHashSet;
+use slice_group_by::GroupByMut;
 use smallvec::{smallvec, SmallVec};
 use std::collections::{HashSet, VecDeque};
 
@@ -1184,7 +1185,7 @@ impl<'a, F: Function> Env<'a, F> {
                 // Find groups of uses that occur in at the same program point.
                 for uses in self.ranges[range.index()]
                     .uses
-                    .group_by_mut_(|a, b| a.pos == b.pos)
+                    .linear_group_by_key_mut(|u| u.pos)
                 {
                     if uses.len() < 2 {
                         continue;
@@ -1254,21 +1255,23 @@ impl<'a, F: Function> Env<'a, F> {
                             // FixedStack is incompatible if there are any
                             // Reg/FixedReg constraints. FixedReg is
                             // incompatible if there already is a different
-                            // FixedReg constraint.
+                            // FixedReg constraint. If either condition is true,
+                            // we edit the constraint below; otherwise, we can
+                            // skip this edit.
                             if !(requires_reg && self.pregs[preg.index()].is_stack)
                                 && *first_preg.get_or_insert(preg) == preg
                             {
                                 continue;
                             }
 
-                            log::trace!(" -> duplicate; switching to constraint Reg");
+                            log::trace!(" -> duplicate; switching to constraint Any");
-                            self.multi_fixed_reg_fixups.push((
+                            self.multi_fixed_reg_fixups.push(MultiFixedRegFixup {
-                                u.pos,
+                                pos: u.pos,
-                                source_slot,
+                                from_slot: source_slot,
-                                preg_idx,
+                                to_slot: u.slot,
-                                vreg_idx,
+                                to_preg: preg_idx,
-                                u.slot,
+                                vreg: vreg_idx,
-                            ));
+                            });
                             u.operand = Operand::new(
                                 u.operand.vreg(),
                                 OperandConstraint::Any,

src/ion/merge.rs

@@ -13,8 +13,8 @@
 //! Bundle merging.
 
 use super::{
-    Env, LiveBundleIndex, LiveRangeIndex, LiveRangeKey, Requirement, SpillSet, SpillSetIndex,
+    Env, LiveBundleIndex, LiveRangeIndex, LiveRangeKey, SpillSet, SpillSetIndex, SpillSlotIndex,
-    SpillSlotIndex, VRegIndex,
+    VRegIndex,
 };
 use crate::{Function, Inst, OperandConstraint, PReg};
 use smallvec::smallvec;
@@ -99,11 +99,7 @@ impl<'a, F: Function> Env<'a, F> {
             || self.bundles[to.index()].cached_stack()
             || self.bundles[to.index()].cached_fixed()
         {
-            let req = self
+            if self.merge_bundle_requirements(from, to).is_err() {
-                .compute_requirement(from)
-                .0
-                .merge(self.compute_requirement(to).0);
-            if req == Requirement::Conflict {
                 log::trace!(" -> conflicting requirements; aborting merge");
                 return false;
             }

src/ion/moves.rs

@@ -695,29 +695,27 @@ impl<'a, F: Function> Env<'a, F> {
         }
 
         // Handle multi-fixed-reg constraints by copying.
-        for (progpoint, from_slot, to_preg, to_vreg, to_slot) in
+        for fixup in std::mem::replace(&mut self.multi_fixed_reg_fixups, vec![]) {
-            std::mem::replace(&mut self.multi_fixed_reg_fixups, vec![])
+            let from_alloc = self.get_alloc(fixup.pos.inst(), fixup.from_slot as usize);
-        {
+            let to_alloc = Allocation::reg(self.pregs[fixup.to_preg.index()].reg);
-            let from_alloc = self.get_alloc(progpoint.inst(), from_slot as usize);
-            let to_alloc = Allocation::reg(self.pregs[to_preg.index()].reg);
             log::trace!(
                 "multi-fixed-move constraint at {:?} from {} to {} for v{}",
-                progpoint,
+                fixup.pos,
                 from_alloc,
                 to_alloc,
-                to_vreg.index(),
+                fixup.vreg.index(),
             );
             self.insert_move(
-                progpoint,
+                fixup.pos,
                 InsertMovePrio::MultiFixedReg,
                 from_alloc,
                 to_alloc,
-                Some(self.vreg_regs[to_vreg.index()]),
+                Some(self.vreg_regs[fixup.vreg.index()]),
             );
             self.set_alloc(
-                progpoint.inst(),
+                fixup.pos.inst(),
-                to_slot as usize,
+                fixup.to_slot as usize,
-                Allocation::reg(self.pregs[to_preg.index()].reg),
+                Allocation::reg(self.pregs[fixup.to_preg.index()].reg),
             );
         }
 

src/ion/process.rs

@@ -18,10 +18,13 @@ use super::{
     Requirement, SpillWeight, UseList,
 };
 use crate::{
-    ion::data_structures::{
+    ion::{
+        data_structures::{
             CodeRange, BUNDLE_MAX_NORMAL_SPILL_WEIGHT, MINIMAL_BUNDLE_SPILL_WEIGHT,
             MINIMAL_FIXED_BUNDLE_SPILL_WEIGHT,
         },
+        requirement::RequirementConflictAt,
+    },
     Allocation, Function, Inst, InstPosition, OperandConstraint, OperandKind, PReg, ProgPoint,
     RegAllocError,
 };
@@ -729,7 +732,6 @@ impl<'a, F: Function> Env<'a, F> {
         reg_hint: PReg,
     ) -> Result<(), RegAllocError> {
         let class = self.spillsets[self.bundles[bundle.index()].spillset.index()].class;
-        let (req, split_point) = self.compute_requirement(bundle);
         // Grab a hint from either the queue or our spillset, if any.
         let mut hint_reg = if reg_hint != PReg::invalid() {
             reg_hint
@@ -741,7 +743,9 @@ impl<'a, F: Function> Env<'a, F> {
         }
         log::trace!("process_bundle: bundle {:?} hint {:?}", bundle, hint_reg,);
 
-        if let Requirement::Conflict = req {
+        let req = match self.compute_requirement(bundle) {
+            Ok(req) => req,
+            Err(RequirementConflictAt(split_point)) => {
                 // We have to split right away. We'll find a point to
                 // split that would allow at least the first half of the
                 // split to be conflict-free.
@@ -756,12 +760,13 @@ impl<'a, F: Function> Env<'a, F> {
                 );
                 return Ok(());
             }
+        };
 
         // If no requirement at all (because no uses), and *if* a
         // spill bundle is already present, then move the LRs over to
         // the spill bundle right away.
         match req {
-            Requirement::Unknown | Requirement::Any => {
+            Requirement::Any => {
                 if let Some(spill) =
                     self.get_or_create_spill_bundle(bundle, /* create_if_absent = */ false)
                 {
@@ -794,14 +799,10 @@ impl<'a, F: Function> Env<'a, F> {
                     return Ok(());
                 }
 
-                Requirement::Any | Requirement::Unknown => {
+                Requirement::Any => {
                     self.spilled_bundles.push(bundle);
                     return Ok(());
                 }
-
-                Requirement::Conflict => {
-                    unreachable!()
-                }
             };
             // Scan all pregs, or the one fixed preg, and attempt to allocate.
 

src/ion/requirement.rs

@@ -15,32 +15,36 @@
 use super::{Env, LiveBundleIndex};
 use crate::{Function, Operand, OperandConstraint, PReg, ProgPoint};
 
+pub struct RequirementConflict;
+
+pub struct RequirementConflictAt(pub ProgPoint);
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum Requirement {
-    Unknown,
     FixedReg(PReg),
     FixedStack(PReg),
     Register,
     Stack,
     Any,
-    Conflict,
 }
 impl Requirement {
     #[inline(always)]
-    pub fn merge(self, other: Requirement) -> Requirement {
+    pub fn merge(self, other: Requirement) -> Result<Requirement, RequirementConflict> {
         match (self, other) {
-            (Requirement::Unknown, other) | (other, Requirement::Unknown) => other,
+            (other, Requirement::Any) | (Requirement::Any, other) => Ok(other),
-            (Requirement::Conflict, _) | (_, Requirement::Conflict) => Requirement::Conflict,
+            (Requirement::Register, Requirement::Register) => Ok(self),
-            (other, Requirement::Any) | (Requirement::Any, other) => other,
+            (Requirement::Stack, Requirement::Stack) => Ok(self),
-            (Requirement::Register, Requirement::Register) => self,
-            (Requirement::Stack, Requirement::Stack) => self,
             (Requirement::Register, Requirement::FixedReg(preg))
-            | (Requirement::FixedReg(preg), Requirement::Register) => Requirement::FixedReg(preg),
+            | (Requirement::FixedReg(preg), Requirement::Register) => {
+                Ok(Requirement::FixedReg(preg))
+            }
             (Requirement::Stack, Requirement::FixedStack(preg))
-            | (Requirement::FixedStack(preg), Requirement::Stack) => Requirement::FixedStack(preg),
+            | (Requirement::FixedStack(preg), Requirement::Stack) => {
-            (Requirement::FixedReg(a), Requirement::FixedReg(b)) if a == b => self,
+                Ok(Requirement::FixedStack(preg))
-            (Requirement::FixedStack(a), Requirement::FixedStack(b)) if a == b => self,
+            }
-            _ => Requirement::Conflict,
+            (Requirement::FixedReg(a), Requirement::FixedReg(b)) if a == b => Ok(self),
+            (Requirement::FixedStack(a), Requirement::FixedStack(b)) if a == b => Ok(self),
+            _ => Err(RequirementConflict),
         }
     }
 }
@@ -58,12 +62,15 @@ impl<'a, F: Function> Env<'a, F> {
             }
             OperandConstraint::Reg | OperandConstraint::Reuse(_) => Requirement::Register,
             OperandConstraint::Stack => Requirement::Stack,
-            _ => Requirement::Any,
+            OperandConstraint::Any => Requirement::Any,
         }
     }
 
-    pub fn compute_requirement(&self, bundle: LiveBundleIndex) -> (Requirement, ProgPoint) {
+    pub fn compute_requirement(
-        let mut req = Requirement::Unknown;
+        &self,
+        bundle: LiveBundleIndex,
+    ) -> Result<Requirement, RequirementConflictAt> {
+        let mut req = Requirement::Any;
         log::trace!("compute_requirement: {:?}", bundle);
         let ranges = &self.bundles[bundle.index()].ranges;
         for entry in ranges {
@@ -71,14 +78,28 @@ impl<'a, F: Function> Env<'a, F> {
             for u in &self.ranges[entry.index.index()].uses {
                 log::trace!("  -> use {:?}", u);
                 let r = self.requirement_from_operand(u.operand);
-                req = req.merge(r);
+                req = req.merge(r).map_err(|_| {
+                    log::trace!("     -> conflict");
+                    RequirementConflictAt(u.pos)
+                })?;
                 log::trace!("     -> req {:?}", req);
-                if req == Requirement::Conflict {
-                    return (req, u.pos);
-                }
             }
         }
         log::trace!(" -> final: {:?}", req);
-        (req, ranges.first().unwrap().range.from)
+        Ok(req)
+    }
+
+    pub fn merge_bundle_requirements(
+        &self,
+        a: LiveBundleIndex,
+        b: LiveBundleIndex,
+    ) -> Result<Requirement, RequirementConflict> {
+        let req_a = self
+            .compute_requirement(a)
+            .map_err(|_| RequirementConflict)?;
+        let req_b = self
+            .compute_requirement(b)
+            .map_err(|_| RequirementConflict)?;
+        req_a.merge(req_b)
     }
 }

src/lib.rs

@@ -19,7 +19,6 @@ pub(crate) mod ion;
 pub(crate) mod moves;
 pub(crate) mod postorder;
 pub(crate) mod ssa;
-pub(crate) mod util;
 
 #[macro_use]
 mod index;

src/util.rs

@@ -1,224 +0,0 @@
-// This file contains a port of the unstable GroupBy slice iterator from the
-// Rust standard library. The methods have a trailing underscore to avoid
-// naming conflicts with the standard library.
-
-use std::fmt;
-use std::iter::FusedIterator;
-use std::mem;
-
-pub trait SliceGroupBy<T> {
-    /// Returns an iterator over the slice producing non-overlapping runs
-    /// of elements using the predicate to separate them.
-    ///
-    /// The predicate is called on two elements following themselves,
-    /// it means the predicate is called on `slice[0]` and `slice[1]`
-    /// then on `slice[1]` and `slice[2]` and so on.
-    fn group_by_<F>(&self, pred: F) -> GroupBy<'_, T, F>
-    where
-        F: FnMut(&T, &T) -> bool;
-
-    /// Returns an iterator over the slice producing non-overlapping mutable
-    /// runs of elements using the predicate to separate them.
-    ///
-    /// The predicate is called on two elements following themselves,
-    /// it means the predicate is called on `slice[0]` and `slice[1]`
-    /// then on `slice[1]` and `slice[2]` and so on.
-    fn group_by_mut_<F>(&mut self, pred: F) -> GroupByMut<'_, T, F>
-    where
-        F: FnMut(&T, &T) -> bool;
-}
-
-impl<T> SliceGroupBy<T> for [T] {
-    fn group_by_<F>(&self, pred: F) -> GroupBy<'_, T, F>
-    where
-        F: FnMut(&T, &T) -> bool,
-    {
-        GroupBy::new(self, pred)
-    }
-
-    fn group_by_mut_<F>(&mut self, pred: F) -> GroupByMut<'_, T, F>
-    where
-        F: FnMut(&T, &T) -> bool,
-    {
-        GroupByMut::new(self, pred)
-    }
-}
-
-/// An iterator over slice in (non-overlapping) chunks separated by a predicate.
-pub struct GroupBy<'a, T: 'a, P> {
-    slice: &'a [T],
-    predicate: P,
-}
-
-impl<'a, T: 'a, P> GroupBy<'a, T, P> {
-    pub(super) fn new(slice: &'a [T], predicate: P) -> Self {
-        GroupBy { slice, predicate }
-    }
-}
-
-impl<'a, T: 'a, P> Iterator for GroupBy<'a, T, P>
-where
-    P: FnMut(&T, &T) -> bool,
-{
-    type Item = &'a [T];
-
-    #[inline]
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.slice.is_empty() {
-            None
-        } else {
-            let mut len = 1;
-            let mut iter = self.slice.windows(2);
-            while let Some([l, r]) = iter.next() {
-                if (self.predicate)(l, r) {
-                    len += 1
-                } else {
-                    break;
-                }
-            }
-            let (head, tail) = self.slice.split_at(len);
-            self.slice = tail;
-            Some(head)
-        }
-    }
-
-    #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        if self.slice.is_empty() {
-            (0, Some(0))
-        } else {
-            (1, Some(self.slice.len()))
-        }
-    }
-
-    #[inline]
-    fn last(mut self) -> Option<Self::Item> {
-        self.next_back()
-    }
-}
-
-impl<'a, T: 'a, P> DoubleEndedIterator for GroupBy<'a, T, P>
-where
-    P: FnMut(&T, &T) -> bool,
-{
-    #[inline]
-    fn next_back(&mut self) -> Option<Self::Item> {
-        if self.slice.is_empty() {
-            None
-        } else {
-            let mut len = 1;
-            let mut iter = self.slice.windows(2);
-            while let Some([l, r]) = iter.next_back() {
-                if (self.predicate)(l, r) {
-                    len += 1
-                } else {
-                    break;
-                }
-            }
-            let (head, tail) = self.slice.split_at(self.slice.len() - len);
-            self.slice = head;
-            Some(tail)
-        }
-    }
-}
-
-impl<'a, T: 'a, P> FusedIterator for GroupBy<'a, T, P> where P: FnMut(&T, &T) -> bool {}
-
-impl<'a, T: 'a + fmt::Debug, P> fmt::Debug for GroupBy<'a, T, P> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("GroupBy")
-            .field("slice", &self.slice)
-            .finish()
-    }
-}
-
-/// An iterator over slice in (non-overlapping) mutable chunks separated
-/// by a predicate.
-pub struct GroupByMut<'a, T: 'a, P> {
-    slice: &'a mut [T],
-    predicate: P,
-}
-
-impl<'a, T: 'a, P> GroupByMut<'a, T, P> {
-    pub(super) fn new(slice: &'a mut [T], predicate: P) -> Self {
-        GroupByMut { slice, predicate }
-    }
-}
-
-impl<'a, T: 'a, P> Iterator for GroupByMut<'a, T, P>
-where
-    P: FnMut(&T, &T) -> bool,
-{
-    type Item = &'a mut [T];
-
-    #[inline]
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.slice.is_empty() {
-            None
-        } else {
-            let mut len = 1;
-            let mut iter = self.slice.windows(2);
-            while let Some([l, r]) = iter.next() {
-                if (self.predicate)(l, r) {
-                    len += 1
-                } else {
-                    break;
-                }
-            }
-            let slice = mem::take(&mut self.slice);
-            let (head, tail) = slice.split_at_mut(len);
-            self.slice = tail;
-            Some(head)
-        }
-    }
-
-    #[inline]
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        if self.slice.is_empty() {
-            (0, Some(0))
-        } else {
-            (1, Some(self.slice.len()))
-        }
-    }
-
-    #[inline]
-    fn last(mut self) -> Option<Self::Item> {
-        self.next_back()
-    }
-}
-
-impl<'a, T: 'a, P> DoubleEndedIterator for GroupByMut<'a, T, P>
-where
-    P: FnMut(&T, &T) -> bool,
-{
-    #[inline]
-    fn next_back(&mut self) -> Option<Self::Item> {
-        if self.slice.is_empty() {
-            None
-        } else {
-            let mut len = 1;
-            let mut iter = self.slice.windows(2);
-            while let Some([l, r]) = iter.next_back() {
-                if (self.predicate)(l, r) {
-                    len += 1
-                } else {
-                    break;
-                }
-            }
-            let slice = mem::take(&mut self.slice);
-            let (head, tail) = slice.split_at_mut(slice.len() - len);
-            self.slice = head;
-            Some(tail)
-        }
-    }
-}
-
-impl<'a, T: 'a, P> FusedIterator for GroupByMut<'a, T, P> where P: FnMut(&T, &T) -> bool {}
-
-impl<'a, T: 'a + fmt::Debug, P> fmt::Debug for GroupByMut<'a, T, P> {
-    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_struct("GroupByMut")
-            .field("slice", &self.slice)
-            .finish()
-    }
-}