wip

[bigint-presentation-code.git] / register_allocator / src / state.rs

diff --git a/register_allocator/src/state.rs b/register_allocator/src/state.rs
index 103233a9991f50b0783664ce581232580727c0cb..684f262ca096dcdd074c32d0a3f7747ca177054b 100644 (file)
--- a/register_allocator/src/state.rs
+++ b/register_allocator/src/state.rs
@@ -1,8 +1,92 @@
-use crate::{function::Function, interned::GlobalState, live_range::LiveRange};
+use crate::{
+    function::Function,
+    index::LiveRangeIdx,
+    interned::GlobalState,
+    live_range::{LiveRange, ProgRange},
+    loc::{Loc, SubLoc},
+};
+use hashbrown::HashMap;
+use std::{collections::BTreeMap, rc::Rc};
 
 #[derive(Debug)]
 pub struct AllocatorState<'a> {
-    global_state: &'a GlobalState,
+    global_state: &'a Rc<GlobalState>,
     func: &'a Function,
     live_ranges: Vec<LiveRange>,
+    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);
+    }
 }