From d7b1cb21ec4027e219fcf41bff5b1296fd8dd124 Mon Sep 17 00:00:00 2001 From: Michael Weigelt Date: Fri, 5 Jun 2026 19:28:44 +0000 Subject: [PATCH 1/2] replace by sorted vec --- src/ion/data_structures.rs | 73 ++++++++++++++++++++++++-- src/ion/liveranges.rs | 1 - src/ion/mod.rs | 2 +- src/ion/moves.rs | 1 - src/ion/process.rs | 102 ++++++++++++++++++------------------- 5 files changed, 121 insertions(+), 58 deletions(-) diff --git a/src/ion/data_structures.rs b/src/ion/data_structures.rs index 41b8ad30..69f23d08 100644 --- a/src/ion/data_structures.rs +++ b/src/ion/data_structures.rs @@ -648,9 +648,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, + 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 { + 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 { + 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 +803,7 @@ impl PrioQueue { impl LiveRangeSet { pub(crate) fn new() -> Self { - Self { - btree: BTreeMap::default(), - } + Self { items: Vec::new() } } } diff --git a/src/ion/liveranges.rs b/src/ion/liveranges.rs index 449f2626..c92f1101 100644 --- a/src/ion/liveranges.rs +++ b/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()); } diff --git a/src/ion/mod.rs b/src/ion/mod.rs index 99181191..898567a3 100644 --- a/src/ion/mod.rs +++ b/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(); diff --git a/src/ion/moves.rs b/src/ion/moves.rs index ddabd9d8..40da10d3 100644 --- a/src/ion/moves.rs +++ b/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); diff --git a/src/ion/process.rs b/src/ion/process.rs index 35f8bb0c..2ec708d2 100644 --- a/src/ion/process.rs +++ b/src/ion/process.rs @@ -64,17 +64,17 @@ 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; @@ -82,17 +82,25 @@ impl<'a, F: Function> Env<'a, F> { 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 = None; let mut last_conflict_bundle: Option = 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,7 +203,7 @@ 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); @@ -208,7 +211,6 @@ impl<'a, F: Function> Env<'a, F> { 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 { From 98504a8dba3aa0c10f29e9cac07720dddbbae694 Mon Sep 17 00:00:00 2001 From: Michael Weigelt Date: Tue, 9 Jun 2026 13:32:41 +0000 Subject: [PATCH 2/2] u64 --- src/ion/data_structures.rs | 27 +++++++++++---------------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/src/ion/data_structures.rs b/src/ion/data_structures.rs index 69f23d08..39af60a1 100644 --- a/src/ion/data_structures.rs +++ b/src/ion/data_structures.rs @@ -447,32 +447,27 @@ impl core::ops::IndexMut 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` 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` +/// 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, - generation: u32, + stamps: Vec, + 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.