register_allocator/src/state.rs

   1 use crate::{
   2     function::Function,
   3     index::LiveRangeIdx,
   4     interned::GlobalState,
   5     live_range::{Lexicographic, LiveRange, ProgRange},
   6     loc::{Loc, SubLoc},
   7 };
   8 use hashbrown::HashMap;
   9 use std::{collections::BTreeMap, rc::Rc};
  10
  11 #[derive(Debug)]
  12 pub struct AllocatorState<'a> {
  13     global_state: &'a Rc<GlobalState>,
  14     func: &'a Function,
  15     live_ranges: Vec<LiveRange>,
  16     allocations: HashMap<SubLoc, BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx>>,
  17 }
  18
  19 impl<'a> AllocatorState<'a> {
  20     pub fn new(global_state: &'a Rc<GlobalState>, func: &'a Function) -> Self {
  21         Self {
  22             global_state,
  23             func,
  24             live_ranges: vec![],
  25             allocations: HashMap::new(),
  26         }
  27     }
  28     pub fn global_state(&self) -> &'a Rc<GlobalState> {
  29         self.global_state
  30     }
  31     pub fn func(&self) -> &'a Function {
  32         self.func
  33     }
  34     pub fn live_ranges(&self) -> &[LiveRange] {
  35         &self.live_ranges
  36     }
  37     pub fn allocations(
  38         &self,
  39         sub_loc: SubLoc,
  40     ) -> &BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx> {
  41         const EMPTY: &'static BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx> = &BTreeMap::new();
  42         self.allocations.get(&sub_loc).unwrap_or(EMPTY)
  43     }
  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);
  51                 }
  52             }
  53         }
  54         old_allocation
  55     }
  56     pub fn add_allocations_for(
  57         &mut self,
  58         live_range_idx: LiveRangeIdx,
  59         new_allocation: Option<Loc>,
  60     ) {
  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() {
  66                 self.allocations
  67                     .entry(sub_loc)
  68                     .or_default()
  69                     .insert(live_range.range, live_range_idx);
  70             }
  71         }
  72     }
  73     pub fn replace_allocation(
  74         &mut self,
  75         live_range_idx: LiveRangeIdx,
  76         new_allocation: Option<Loc>,
  77     ) -> Option<Loc> {
  78         let old_allocation = self.remove_allocations_for(live_range_idx);
  79         self.add_allocations_for(live_range_idx, new_allocation);
  80         old_allocation
  81     }
  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);
  87         live_range_idx
  88     }
  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);
  94     }
  95 }