+ allocations: HashMap<SubLoc, BTreeMap<ProgRange, LiveRangeIdx>>,
+}
+
+impl<'a> AllocatorState<'a> {
+ pub fn new(global_state: &'a Rc<GlobalState>, func: &'a Function) -> Self {
+ Self {
+ global_state,
+ func,
+ live_ranges: vec![],
+ allocations: HashMap::new(),
+ }
+ }
+ pub fn global_state(&self) -> &'a Rc<GlobalState> {
+ self.global_state
+ }
+ pub fn func(&self) -> &'a Function {
+ self.func
+ }
+ pub fn live_ranges(&self) -> &[LiveRange] {
+ &self.live_ranges
+ }
+ pub fn allocations(&self, sub_loc: SubLoc) -> &BTreeMap<ProgRange, LiveRangeIdx> {
+ const EMPTY: &'static BTreeMap<ProgRange, LiveRangeIdx> = &BTreeMap::new();
+ self.allocations.get(&sub_loc).unwrap_or(EMPTY)
+ }
+ pub fn remove_allocations_for(&mut self, live_range_idx: LiveRangeIdx) -> Option<Loc> {
+ let live_range = &mut self.live_ranges[live_range_idx];
+ let old_allocation = live_range.allocation.take();
+ if let Some(old_allocation) = old_allocation {
+ for sub_loc in old_allocation.sub_locs() {
+ if let Some(allocations) = self.allocations.get_mut(&sub_loc) {
+ allocations.remove(&live_range.range);
+ }
+ }
+ }
+ old_allocation
+ }
+ pub fn add_allocations_for(
+ &mut self,
+ live_range_idx: LiveRangeIdx,
+ new_allocation: Option<Loc>,
+ ) {
+ let live_range = &mut self.live_ranges[live_range_idx];
+ assert!(live_range.allocation.is_none());
+ live_range.allocation = new_allocation;
+ if let Some(new_allocation) = new_allocation {
+ for sub_loc in new_allocation.sub_locs() {
+ self.allocations
+ .entry(sub_loc)
+ .or_default()
+ .insert(live_range.range, live_range_idx);
+ }
+ }
+ }
+ pub fn replace_allocation(
+ &mut self,
+ live_range_idx: LiveRangeIdx,
+ new_allocation: Option<Loc>,
+ ) -> Option<Loc> {
+ let old_allocation = self.remove_allocations_for(live_range_idx);
+ self.add_allocations_for(live_range_idx, new_allocation);
+ old_allocation
+ }
+ pub fn add_live_range(&mut self, mut live_range: LiveRange) -> LiveRangeIdx {
+ let allocation = live_range.allocation.take();
+ let live_range_idx = LiveRangeIdx::new(self.live_ranges.len());
+ self.live_ranges.push(live_range);
+ self.add_allocations_for(live_range_idx, allocation);
+ live_range_idx
+ }
+ pub fn set_live_range(&mut self, live_range_idx: LiveRangeIdx, mut new_live_range: LiveRange) {
+ let allocation = new_live_range.allocation.take();
+ self.remove_allocations_for(live_range_idx);
+ self.live_ranges[live_range_idx] = new_live_range;
+ self.add_allocations_for(live_range_idx, allocation);
+ }