register_allocator/src/state.rs

   1 use crate::{
   2     function::Function,
   3     index::LiveRangeIdx,
   4     interned::GlobalState,
   5     live_range::{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, 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(&self, sub_loc: SubLoc) -> &BTreeMap<ProgRange, LiveRangeIdx> {
  38         const EMPTY: &'static BTreeMap<ProgRange, LiveRangeIdx> = &BTreeMap::new();
  39         self.allocations.get(&sub_loc).unwrap_or(EMPTY)
  40     }
  41     pub fn remove_allocations_for(&mut self, live_range_idx: LiveRangeIdx) -> Option<Loc> {
  42         let live_range = &mut self.live_ranges[live_range_idx];
  43         let old_allocation = live_range.allocation.take();
  44         if let Some(old_allocation) = old_allocation {
  45             for sub_loc in old_allocation.sub_locs() {
  46                 if let Some(allocations) = self.allocations.get_mut(&sub_loc) {
  47                     allocations.remove(&live_range.range);
  48                 }
  49             }
  50         }
  51         old_allocation
  52     }
  53     pub fn add_allocations_for(
  54         &mut self,
  55         live_range_idx: LiveRangeIdx,
  56         new_allocation: Option<Loc>,
  57     ) {
  58         let live_range = &mut self.live_ranges[live_range_idx];
  59         assert!(live_range.allocation.is_none());
  60         live_range.allocation = new_allocation;
  61         if let Some(new_allocation) = new_allocation {
  62             for sub_loc in new_allocation.sub_locs() {
  63                 self.allocations
  64                     .entry(sub_loc)
  65                     .or_default()
  66                     .insert(live_range.range, live_range_idx);
  67             }
  68         }
  69     }
  70     pub fn replace_allocation(
  71         &mut self,
  72         live_range_idx: LiveRangeIdx,
  73         new_allocation: Option<Loc>,
  74     ) -> Option<Loc> {
  75         let old_allocation = self.remove_allocations_for(live_range_idx);
  76         self.add_allocations_for(live_range_idx, new_allocation);
  77         old_allocation
  78     }
  79     pub fn add_live_range(&mut self, mut live_range: LiveRange) -> LiveRangeIdx {
  80         let allocation = live_range.allocation.take();
  81         let live_range_idx = LiveRangeIdx::new(self.live_ranges.len());
  82         self.live_ranges.push(live_range);
  83         self.add_allocations_for(live_range_idx, allocation);
  84         live_range_idx
  85     }
  86     pub fn set_live_range(&mut self, live_range_idx: LiveRangeIdx, mut new_live_range: LiveRange) {
  87         let allocation = new_live_range.allocation.take();
  88         self.remove_allocations_for(live_range_idx);
  89         self.live_ranges[live_range_idx] = new_live_range;
  90         self.add_allocations_for(live_range_idx, allocation);
  91     }
  92 }