perf: Replace BTreeMap with sorted vec

src/ion/data_structures.rs

@@ -447,32 +447,27 @@ impl core::ops::IndexMut<VReg> for VRegs {
 /// A dedup set of `LiveBundleIndex` values that avoids hashing.
 ///
 /// Each bundle index is mapped to a slot in a dense array holding the
-/// generation at which it was last inserted. `clear` simply bumps the
-/// current generation (O(1) in the common case), and `insert` is a
-/// single bounds-checked compare-and-store. This is a drop-in
-/// replacement for the previous `FxHashSet<LiveBundleIndex>` that is
-/// much cheaper when a bundle conflicts with many others (e.g. a
-/// function with many locals).
+/// generation at which it was last inserted. `clear` bumps the current
+/// generation and `insert` updates the stamp for the given bundle and
+/// resizes the array if necessary.
+/// This is a drop-in replacement for the previous `FxHashSet<LiveBundleIndex>`
+/// that is much cheaper when a bundle conflicts with many others (e.g. for
+/// functions with many locals).
 #[derive(Clone, Debug, Default)]
 pub struct ConflictSet {
     // Generation at which each bundle was last inserted. The value `0`
     // means "never inserted"; `generation` is therefore always >= 1
     // after the first `clear`.
-    stamps: Vec<u32>,
-    generation: u32,
+    stamps: Vec<u64>,
+    generation: u64,
 }
 
 impl ConflictSet {
-    /// Empty the set. O(1) except on the rare generation wraparound.
+    /// Empty the set by bumping the generation.
+    /// Every value below the new generation is now stale.
     #[inline]
     pub fn clear(&mut self) {
-        self.generation = self.generation.wrapping_add(1);
-        if self.generation == 0 {
-            // Wrapped around; reset stamps so stale entries don't read
-            // as present, and skip the reserved `0` generation.
-            self.stamps.iter_mut().for_each(|s| *s = 0);
-            self.generation = 1;
-        }
+        self.generation += 1;
     }
 
     /// Insert `bundle`. Returns `true` if it was not already present.
@@ -648,9 +643,76 @@ pub struct PrioQueueEntry {
     pub hint: PReg,
 }
 
+/// The set of live ranges currently allocated to a single PReg.
+///
+/// Entries are kept sorted and non-overlapping by `LiveRangeKey`. This
+/// is backed by a flat `Vec` rather than a `BTreeMap` because the hot
+/// path (`try_to_allocate_bundle_to_reg`) is forward iteration over a
+/// sub-range while scanning for overlaps; a contiguous slice walk is
+/// dramatically cheaper and more cache-friendly than the std
+/// `BTreeMap` iterator (which does per-element handle juggling). Lookup
+/// is a binary search, and insert/remove shift the tail.
+/// The trade-off is that mutations are O(log n) rather than O(1).
 #[derive(Clone, Debug)]
 pub struct LiveRangeSet {
-    pub btree: BTreeMap<LiveRangeKey, LiveRangeIndex>,
+    pub items: Vec<(LiveRangeKey, LiveRangeIndex)>,
+}
+
+impl LiveRangeSet {
+    /// Index of the first entry whose key is not ordered strictly
+    /// before `key`, i.e., the first entry `>= key` under the
+    /// overlap-aware `LiveRangeKey` ordering. Equivalent to the start
+    /// of `BTreeMap::range(key..)`.
+    #[inline(always)]
+    fn lower_bound(&self, key: &LiveRangeKey) -> usize {
+        self.items.partition_point(|(k, _)| k < key)
+    }
+
+    /// The entries starting at the first one that is `>= key`, as a
+    /// contiguous slice. Equivalent to `BTreeMap::range(key..)`.
+    #[inline(always)]
+    pub fn range_from(&self, key: LiveRangeKey) -> &[(LiveRangeKey, LiveRangeIndex)] {
+        let start = self.lower_bound(&key);
+        &self.items[start..]
+    }
+
+    /// Insert `(key, value)`, keeping the set sorted. Returns the
+    /// previous value if an entry with an equal (overlapping) key was
+    /// already present, mirroring `BTreeMap::insert`.
+    #[inline]
+    pub fn insert(&mut self, key: LiveRangeKey, value: LiveRangeIndex) -> Option<LiveRangeIndex> {
+        let idx = self.lower_bound(&key);
+        if idx < self.items.len() && self.items[idx].0 == key {
+            Some(core::mem::replace(&mut self.items[idx].1, value))
+        } else {
+            self.items.insert(idx, (key, value));
+            None
+        }
+    }
+
+    /// Remove the entry whose key equals (overlaps) `key`. Returns its
+    /// value if present, mirroring `BTreeMap::remove`.
+    #[inline]
+    pub fn remove(&mut self, key: &LiveRangeKey) -> Option<LiveRangeIndex> {
+        let idx = self.lower_bound(key);
+        if idx < self.items.len() && self.items[idx].0 == *key {
+            Some(self.items.remove(idx).1)
+        } else {
+            None
+        }
+    }
+
+    /// Whether any entry's key equals (overlaps) `key`.
+    #[inline]
+    pub fn contains_key(&self, key: &LiveRangeKey) -> bool {
+        let idx = self.lower_bound(key);
+        idx < self.items.len() && self.items[idx].0 == *key
+    }
+
+    #[inline(always)]
+    pub fn clear(&mut self) {
+        self.items.clear();
+    }
 }
 
 #[derive(Clone, Copy, Debug)]
@@ -736,9 +798,7 @@ impl PrioQueue {
 
 impl LiveRangeSet {
     pub(crate) fn new() -> Self {
-        Self {
-            btree: BTreeMap::default(),
-        }
+        Self { items: Vec::new() }
     }
 }
 

src/ion/liveranges.rs

@@ -264,7 +264,6 @@ impl<'a, F: Function> Env<'a, F> {
         let preg_idx = PRegIndex::new(reg.index());
         let res = self.pregs[preg_idx.index()]
             .allocations
-            .btree
             .insert(LiveRangeKey::from_range(&range), LiveRangeIndex::invalid());
         debug_assert!(res.is_none());
     }

src/ion/mod.rs

@@ -50,7 +50,7 @@ impl<'a, F: Function> Env<'a, F> {
         ctx.vregs.preallocate(ninstrs);
         for preg in ctx.pregs.iter_mut() {
             preg.is_stack = false;
-            preg.allocations.btree.clear();
+            preg.allocations.clear();
         }
         ctx.allocation_queue.heap.clear();
         ctx.spilled_bundles.clear();

src/ion/moves.rs

@@ -910,7 +910,6 @@ impl<'a, F: Function> Env<'a, F> {
                     while let Some(preg) = scratch_iter.next() {
                         if !self.pregs[preg.index()]
                             .allocations
-                            .btree
                             .contains_key(&key)
                         {
                             let alloc = Allocation::reg(preg);

src/ion/process.rs

@@ -64,35 +64,43 @@ impl<'a, F: Function> Env<'a, F> {
         conflicts.clear();
         self.ctx.conflict_set.clear();
         let mut max_conflict_weight = 0;
-        // Traverse the BTreeMap in order by requesting the whole
-        // range spanned by the bundle and iterating over that
-        // concurrently with our ranges. Because our ranges are in
-        // order, and the BTreeMap is as well, this allows us to have
-        // an overall O(n log n) + O(b) complexity, where the PReg has
-        // n current ranges and the bundle has b ranges, rather than
-        // O(b * n log n) with the simple probe-for-each-bundle-range
-        // approach.
+        // Traverse the PReg's sorted range list in order, binary
+        // searching to the first entry spanned by the bundle and then
+        // iterating over that slice concurrently with our ranges.
+        // Because both our ranges and the PReg's are in order, this
+        // gives an overall O(n + b) walk (plus an O(log n) seek), where
+        // the PReg has n current ranges and the bundle has b ranges,
+        // rather than O(b * n log n) with the simple
+        // probe-for-each-bundle-range approach.
         //
-        // Note that the comparator function on a CodeRange tests for
-        // *overlap*, so we are checking whether the BTree contains
+        // Note that the comparator function on a LiveRangeKey tests for
+        // *overlap*, so we are checking whether the PReg list contains
         // any preg range that *overlaps* with range `range`, not
         // literally the range `range`.
         let bundle_ranges = &self.ctx.bundles[bundle].ranges;
         let from_key = LiveRangeKey::from_range(&CodeRange {
             from: bundle_ranges.first().unwrap().range.from,
             to: bundle_ranges.first().unwrap().range.from,
         });
-        let mut preg_range_iter = self.ctx.pregs[reg.index()]
-            .allocations
-            .btree
-            .range(from_key..)
-            .peekable();
+        // The PReg's allocated ranges, kept as a sorted, non-overlapping,
+        // contiguous slice. We walk it concurrently with the bundle's
+        // ranges using a monotonically advancing index `pos`; a slice
+        // walk avoids the per-element cost of the std BTreeMap iterator
+        // and is far more cache-friendly.
+        let preg_items = self.ctx.pregs[reg.index()].allocations.items.as_slice();
         trace!(
             "alloc map for {:?} in range {:?}..: {:?}",
             reg,
             from_key,
-            self.ctx.pregs[reg.index()].allocations.btree
+            preg_items
         );
+        // A binary-search re-seek costs ~O(log n), so we only fall back
+        // to one once we would otherwise step over more than that many
+        // entries. Scaling the threshold with the slice size (~log2(n))
+        // keeps the common small case stepping cheaply while bounding
+        // the worst-case skip work on a densely populated PReg.
+        let reseek_threshold = (usize::BITS - preg_items.len().max(1).leading_zeros()) as usize;
+        let mut pos = preg_items.partition_point(|(k, _)| *k < from_key);
         let mut first_conflict: Option<ProgPoint> = None;
         let mut last_conflict_bundle: Option<LiveBundleIndex> = None;
 
@@ -102,62 +110,57 @@ impl<'a, F: Function> Env<'a, F> {
 
             let mut skips = 0;
             'alloc: loop {
-                trace!("  -> PReg range {:?}", preg_range_iter.peek());
-
-                // Advance our BTree traversal until it is >= this bundle
-                // range (i.e., skip PReg allocations in the BTree that
-                // are completely before this bundle range).
-
-                if preg_range_iter.peek().is_some() && *preg_range_iter.peek().unwrap().0 < key {
-                    trace!(
-                        "Skipping PReg range {:?}",
-                        preg_range_iter.peek().unwrap().0
-                    );
-                    preg_range_iter.next();
+                // If there are no more PReg allocations, we're done!
+                let (cur_key, cur_range) = match preg_items.get(pos) {
+                    None => {
+                        trace!(" -> no more PReg allocations; so no conflict possible!");
+                        break 'ranges;
+                    }
+                    Some(&(k, v)) => (k, v),
+                };
+                trace!("  -> PReg range {:?}", (cur_key, cur_range));
+
+                // Advance our traversal until it is >= this bundle range
+                // (i.e., skip PReg allocations that are completely before
+                // this bundle range).
+                if cur_key < key {
+                    trace!("Skipping PReg range {:?}", cur_key);
+                    pos += 1;
                     skips += 1;
-                    if skips >= 16 {
+                    if skips >= reseek_threshold {
                         let from_pos = entry.range.from;
                         let from_key = LiveRangeKey::from_range(&CodeRange {
                             from: from_pos,
                             to: from_pos,
                         });
-                        preg_range_iter = self.ctx.pregs[reg.index()]
-                            .allocations
-                            .btree
-                            .range(from_key..)
-                            .peekable();
+                        pos += preg_items[pos..].partition_point(|(k, _)| *k < from_key);
                         skips = 0;
                     }
                     continue 'alloc;
                 }
                 skips = 0;
 
-                // If there are no more PReg allocations, we're done!
-                if preg_range_iter.peek().is_none() {
-                    trace!(" -> no more PReg allocations; so no conflict possible!");
-                    break 'ranges;
-                }
-
                 // If the current PReg range is beyond this range, there is no conflict; continue.
-                if *preg_range_iter.peek().unwrap().0 > key {
+                if cur_key > key {
                     trace!(
                         " -> next PReg allocation is at {:?}; moving to next VReg range",
-                        preg_range_iter.peek().unwrap().0
+                        cur_key
                     );
                     break 'alloc;
                 }
 
                 // Otherwise, there is a conflict.
-                let preg_key = *preg_range_iter.peek().unwrap().0;
+                let preg_key = cur_key;
                 debug_assert_eq!(preg_key, key); // Assert that this range overlaps.
-                let preg_range = preg_range_iter.next().unwrap().1;
+                let preg_range = cur_range;
+                pos += 1;
 
                 trace!(" -> btree contains range {:?} that overlaps", preg_range);
                 if preg_range.is_valid() {
-                    trace!("   -> from vreg {:?}", self.ctx.ranges[*preg_range].vreg);
+                    trace!("   -> from vreg {:?}", self.ctx.ranges[preg_range].vreg);
                     // range from an allocated bundle: find the bundle and add to
                     // conflicts list.
-                    let conflict_bundle = self.ctx.ranges[*preg_range].bundle;
+                    let conflict_bundle = self.ctx.ranges[preg_range].bundle;
                     trace!("   -> conflict bundle {:?}", conflict_bundle);
                     // Adjacent preg ranges very often belong to the same
                     // bundle; skip the dedup-set lookup on repeats.
@@ -200,15 +203,14 @@ impl<'a, F: Function> Env<'a, F> {
             return AllocRegResult::Conflict(conflicts, first_conflict.unwrap());
         }
 
-        // We can allocate! Add our ranges to the preg's BTree.
+        // We can allocate! Add our ranges to the preg's allocation set.
         let preg = PReg::from_index(reg.index());
         trace!("  -> bundle {:?} assigned to preg {:?}", bundle, preg);
         self.ctx.bundles[bundle].allocation = Allocation::reg(preg);
         for entry in &self.ctx.bundles[bundle].ranges {
             let key = LiveRangeKey::from_range(&entry.range);
             let res = self.ctx.pregs[reg.index()]
                 .allocations
-                .btree
                 .insert(key, entry.index);
 
             // We disallow LR overlap within bundles, so this should never be possible.
@@ -240,7 +242,6 @@ impl<'a, F: Function> Env<'a, F> {
             trace!(" -> removing LR {:?} from reg {:?}", entry.index, preg_idx);
             self.ctx.pregs[preg_idx.index()]
                 .allocations
-                .btree
                 .remove(&LiveRangeKey::from_range(&entry.range));
         }
         let prio = self.ctx.bundles[bundle].prio;
@@ -1236,8 +1237,7 @@ impl<'a, F: Function> Env<'a, F> {
                         });
                         for (key, lr) in self.ctx.pregs[preg.index()]
                             .allocations
-                            .btree
-                            .range(start..)
+                            .range_from(start)
                         {
                             let preg_range = key.to_range();
                             if preg_range.to <= range.from {