reusing memory on more places

src/ion/data_structures.rs

@@ -21,14 +21,63 @@ use crate::{
     define_index, Allocation, Block, Edit, Function, FxHashMap, FxHashSet, MachineEnv, Operand,
     Output, PReg, ProgPoint, RegClass, VReg, VecExt,
 };
-use alloc::collections::BTreeMap;
+//use alloc::collections::BTreeMap;
 use alloc::string::String;
 use alloc::vec::Vec;
 use core::cmp::Ordering;
+use core::convert::identity;
 use core::fmt::Debug;
 use core::ops::{Deref, DerefMut};
 use smallvec::{smallvec, SmallVec};
 
+#[derive(Clone, Debug)]
+pub struct BTreeMap<K, V> {
+    values: Vec<(K, V)>,
+}
+
+impl<K: Eq + Ord + Copy, V> BTreeMap<K, V> {
+    pub fn clear(&mut self) {
+        self.values.clear();
+    }
+
+    pub fn contains_key(&self, key: &K) -> bool {
+        self.values.binary_search_by_key(key, |e| e.0).is_ok()
+    }
+
+    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        match self.values.binary_search_by_key(&key, |e| e.0) {
+            Ok(i) => Some(core::mem::replace(&mut self.values[i], (key, value)).1),
+            Err(i) => {
+                self.values.insert(i, (key, value));
+                None
+            }
+        }
+    }
+
+    pub fn remove(&mut self, key: &K) -> Option<V> {
+        self.values
+            .binary_search_by_key(key, |e| e.0)
+            .ok()
+            .map(|i| self.values.remove(i).1)
+    }
+
+    pub fn range(&self, range: std::ops::RangeFrom<K>) -> impl Iterator<Item = (&K, &V)> {
+        let start = self
+            .values
+            .binary_search_by_key(&range.start, |e| e.0)
+            .unwrap_or_else(identity);
+        self.values[start..].iter().map(|(k, v)| (k, v))
+    }
+}
+
+impl<K, V> Default for BTreeMap<K, V> {
+    fn default() -> Self {
+        Self {
+            values: Default::default(),
+        }
+    }
+}
+
 /// A range from `from` (inclusive) to `to` (exclusive).
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct CodeRange {
@@ -476,6 +525,7 @@ pub struct Ctx {
     // Output:
     pub output: Output,
 
+    pub(crate) scratch_spillset_pool: Vec<SpillSetRanges>,
     pub(crate) scratch_moves: MoveCtx,
     pub(crate) scratch_removed_lrs: FxHashSet<LiveRangeIndex>,
     pub(crate) scratch_removed_lrs_vregs: FxHashSet<VRegIndex>,
@@ -528,19 +578,11 @@ impl<'a, F: Function> Env<'a, F> {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 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: SpillSetRanges,
@@ -668,7 +710,7 @@ impl PrioQueue {
 impl LiveRangeSet {
     pub(crate) fn new() -> Self {
         Self {
-            btree: BTreeMap::new(),
+            btree: BTreeMap::default(),
         }
     }
 }

src/ion/mod.rs

@@ -48,10 +48,14 @@ impl<'a, F: Function> Env<'a, F> {
         ctx.bundles.prepare(ninstrs);
         ctx.spillsets.prepare(ninstrs);
         ctx.vregs.prepare(ninstrs);
-        ctx.pregs.clear();
+        for preg in ctx.pregs.iter_mut() {
+            preg.is_stack = false;
+            preg.allocations.btree.clear();
+        }
         ctx.allocation_queue.heap.clear();
         ctx.spilled_bundles.clear();
-        ctx.spillslots.clear();
+        ctx.scratch_spillset_pool
+            .extend(ctx.spillslots.drain(..).map(|s| s.ranges));
         ctx.slots_by_class = core::array::from_fn(|_| SpillSlotList::default());
         ctx.extra_spillslots_by_class = core::array::from_fn(|_| smallvec![]);
         ctx.preferred_victim_by_class = [PReg::invalid(); 3];

src/ion/process.rs

@@ -187,6 +187,8 @@ impl<'a, F: Function> Env<'a, F> {
             }
         }
 
+        drop(preg_range_iter);
+
         if conflicts.len() > 0 {
             return AllocRegResult::Conflict(conflicts, first_conflict.unwrap());
         }

src/ion/spill.rs

@@ -132,7 +132,7 @@ impl<'a, F: Function> Env<'a, F> {
                 // Allocate a new spillslot.
                 let spillslot = SpillSlotIndex::new(self.ctx.spillslots.len());
                 self.ctx.spillslots.push(SpillSlotData {
-                    ranges: SpillSetRanges::new(),
+                    ranges: self.ctx.scratch_spillset_pool.pop().unwrap_or_default(),
                     alloc: Allocation::none(),
                     slots: self.func.spillslot_size(self.ctx.spillsets[spillset].class) as u32,
                 });