Initial support for overlapping LiveRanges (#122)

Relax the restrictions on liverange overlap for VRegs.

#122

Author: elliottt

src/ion/data_structures.rs

@@ -38,7 +38,7 @@ pub struct CodeRange {
 impl CodeRange {
     #[inline(always)]
     pub fn is_empty(&self) -> bool {
-        self.from == self.to
+        self.from >= self.to
     }
     #[inline(always)]
     pub fn contains(&self, other: &Self) -> bool {
@@ -64,6 +64,15 @@ impl CodeRange {
             to: pos.next(),
         }
     }
+
+    /// Join two [CodeRange] values together, producing a [CodeRange] that includes both.
+    #[inline(always)]
+    pub fn join(&self, other: CodeRange) -> Self {
+        CodeRange {
+            from: self.from.min(other.from),
+            to: self.to.max(other.to),
+        }
+    }
 }
 
 impl core::cmp::PartialOrd for CodeRange {
@@ -153,6 +162,9 @@ impl LiveRange {
     }
     #[inline(always)]
     pub fn uses_spill_weight(&self) -> SpillWeight {
+        // NOTE: the spill weight is technically stored in 29 bits, but we ignore the sign bit as
+        // we will always be dealing with positive values. Thus we mask out the top 3 bits to
+        // ensure that the sign bit is clear, then shift left by only two.
         let bits = (self.uses_spill_weight_and_flags & 0x1fff_ffff) << 2;
         SpillWeight::from_f32(f32::from_bits(bits))
     }
@@ -285,17 +297,22 @@ const fn no_bloat_capacity<T>() -> usize {
 
 #[derive(Clone, Debug)]
 pub struct SpillSet {
-    pub vregs: SmallVec<[VRegIndex; no_bloat_capacity::<VRegIndex>()]>,
     pub slot: SpillSlotIndex,
     pub reg_hint: PReg,
     pub class: RegClass,
     pub spill_bundle: LiveBundleIndex,
     pub required: bool,
     pub size: u8,
     pub splits: u8,
+
+    /// The aggregate [`CodeRange`] of all involved [`LiveRange`]s. The effect of this abstraction
+    /// is that we attempt to allocate one spill slot for the extent of a bundle. For fragmented
+    /// bundles with lots of open space this abstraction is pessimistic, but when bundles are small
+    /// or dense this yields similar results to tracking individual live ranges.
+    pub range: CodeRange,
 }
 
-pub(crate) const MAX_SPLITS_PER_SPILLSET: u8 = 10;
+pub(crate) const MAX_SPLITS_PER_SPILLSET: u8 = 2;
 
 #[derive(Clone, Debug)]
 pub struct VRegData {
@@ -466,9 +483,22 @@ impl<'a, F: Function> Env<'a, F> {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct SpillSetRanges {
+    pub btree: BTreeMap<LiveRangeKey, SpillSetIndex>,
+}
+
+impl SpillSetRanges {
+    pub fn new() -> Self {
+        Self {
+            btree: BTreeMap::new(),
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct SpillSlotData {
-    pub ranges: LiveRangeSet,
+    pub ranges: SpillSetRanges,
     pub slots: u32,
     pub alloc: Allocation,
 }

src/ion/liveranges.rs

@@ -309,10 +309,11 @@ impl<'a, F: Function> Env<'a, F> {
     pub fn add_liverange_to_preg(&mut self, range: CodeRange, reg: PReg) {
         trace!("adding liverange to preg: {:?} to {}", range, reg);
         let preg_idx = PRegIndex::new(reg.index());
-        self.pregs[preg_idx.index()]
+        let res = self.pregs[preg_idx.index()]
             .allocations
             .btree
             .insert(LiveRangeKey::from_range(&range), LiveRangeIndex::invalid());
+        debug_assert!(res.is_none());
     }
 
     pub fn is_live_in(&mut self, block: Block, vreg: VRegIndex) -> bool {

src/ion/merge.rs

@@ -214,16 +214,10 @@ impl<'a, F: Function> Env<'a, F> {
         }
 
         if self.bundles[from.index()].spillset != self.bundles[to.index()].spillset {
-            let from_vregs = core::mem::replace(
-                &mut self.spillsets[self.bundles[from.index()].spillset.index()].vregs,
-                smallvec![],
-            );
-            let to_vregs = &mut self.spillsets[self.bundles[to.index()].spillset.index()].vregs;
-            for vreg in from_vregs {
-                if !to_vregs.contains(&vreg) {
-                    to_vregs.push(vreg);
-                }
-            }
+            // Widen the range for the target spillset to include the one being merged in.
+            let from_range = self.spillsets[self.bundles[from.index()].spillset.index()].range;
+            let to_range = &mut self.spillsets[self.bundles[to.index()].spillset.index()].range;
+            *to_range = to_range.join(from_range);
         }
 
         if self.bundles[from.index()].cached_stack() {
@@ -246,6 +240,8 @@ impl<'a, F: Function> Env<'a, F> {
             }
 
             let bundle = self.create_bundle();
+            let mut range = self.vregs[vreg.index()].ranges.first().unwrap().range;
+
             self.bundles[bundle.index()].ranges = self.vregs[vreg.index()].ranges.clone();
             trace!("vreg v{} gets bundle{}", vreg.index(), bundle.index());
             for entry in &self.bundles[bundle.index()].ranges {
@@ -254,6 +250,7 @@ impl<'a, F: Function> Env<'a, F> {
                     entry.index.index(),
                     entry.range
                 );
+                range = range.join(entry.range);
                 self.ranges[entry.index.index()].bundle = bundle;
             }
 
@@ -291,14 +288,14 @@ impl<'a, F: Function> Env<'a, F> {
             let reg = self.vreg(vreg);
             let size = self.func.spillslot_size(reg.class()) as u8;
             self.spillsets.push(SpillSet {
-                vregs: smallvec![vreg],
                 slot: SpillSlotIndex::invalid(),
                 size,
                 required: false,
                 class: reg.class(),
                 reg_hint: PReg::invalid(),
                 spill_bundle: LiveBundleIndex::invalid(),
                 splits: 0,
+                range,
             });
             self.bundles[bundle.index()].spillset = ssidx;
         }