5 live_range::{Lexicographic, LiveRange, ProgRange},
8 use hashbrown::HashMap;
9 use std::{collections::BTreeMap, rc::Rc};
12 pub struct AllocatorState<'a> {
13 global_state: &'a Rc<GlobalState>,
15 live_ranges: Vec<LiveRange>,
16 allocations: HashMap<SubLoc, BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx>>,
19 impl<'a> AllocatorState<'a> {
20 pub fn new(global_state: &'a Rc<GlobalState>, func: &'a Function) -> Self {
25 allocations: HashMap::new(),
28 pub fn global_state(&self) -> &'a Rc<GlobalState> {
31 pub fn func(&self) -> &'a Function {
34 pub fn live_ranges(&self) -> &[LiveRange] {
40 ) -> &BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx> {
41 const EMPTY: &'static BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx> = &BTreeMap::new();
42 self.allocations.get(&sub_loc).unwrap_or(EMPTY)
44 pub fn remove_allocations_for(&mut self, live_range_idx: LiveRangeIdx) -> Option<Loc> {
45 let live_range = &mut self.live_ranges[live_range_idx];
46 let old_allocation = live_range.allocation.take();
47 if let Some(old_allocation) = old_allocation {
48 for sub_loc in old_allocation.sub_locs() {
49 if let Some(allocations) = self.allocations.get_mut(&sub_loc) {
50 allocations.remove(&live_range.range);
56 pub fn add_allocations_for(
58 live_range_idx: LiveRangeIdx,
59 new_allocation: Option<Loc>,
61 let live_range = &mut self.live_ranges[live_range_idx];
62 assert!(live_range.allocation.is_none());
63 live_range.allocation = new_allocation;
64 if let Some(new_allocation) = new_allocation {
65 for sub_loc in new_allocation.sub_locs() {
69 .insert(live_range.range, live_range_idx);
73 pub fn replace_allocation(
75 live_range_idx: LiveRangeIdx,
76 new_allocation: Option<Loc>,
78 let old_allocation = self.remove_allocations_for(live_range_idx);
79 self.add_allocations_for(live_range_idx, new_allocation);
82 pub fn add_live_range(&mut self, mut live_range: LiveRange) -> LiveRangeIdx {
83 let allocation = live_range.allocation.take();
84 let live_range_idx = LiveRangeIdx::new(self.live_ranges.len());
85 self.live_ranges.push(live_range);
86 self.add_allocations_for(live_range_idx, allocation);
89 pub fn set_live_range(&mut self, live_range_idx: LiveRangeIdx, mut new_live_range: LiveRange) {
90 let allocation = new_live_range.allocation.take();
91 self.remove_allocations_for(live_range_idx);
92 self.live_ranges[live_range_idx] = new_live_range;
93 self.add_allocations_for(live_range_idx, allocation);