From 87a329a262117b402b9c62b55069bcc20f2b0a66 Mon Sep 17 00:00:00 2001
From: Jacob Lifshay <programmerjake@gmail.com>
Date: Thu, 16 Feb 2023 19:59:49 -0800
Subject: [PATCH] add ssa_val_allocations

---
 register_allocator/src/state.rs | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/register_allocator/src/state.rs b/register_allocator/src/state.rs
index acbdc64..3bcb5cd 100644
--- a/register_allocator/src/state.rs
+++ b/register_allocator/src/state.rs
@@ -1,6 +1,6 @@
 use crate::{
     function::Function,
-    index::LiveRangeIdx,
+    index::{Entries, LiveRangeIdx, SSAValIdx},
     interned::GlobalState,
     live_range::{Lexicographic, LiveRange, ProgRange},
     loc::{Loc, SubLoc},
@@ -13,7 +13,8 @@ pub struct AllocatorState<'a> {
     global_state: &'a Rc<GlobalState>,
     func: &'a Function,
     live_ranges: Vec<LiveRange>,
-    allocations: HashMap<SubLoc, BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx>>,
+    sub_loc_allocations: HashMap<SubLoc, BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx>>,
+    ssa_val_allocations: Box<[BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx>]>,
 }
 
 impl<'a> AllocatorState<'a> {
@@ -22,7 +23,8 @@ impl<'a> AllocatorState<'a> {
             global_state,
             func,
             live_ranges: vec![],
-            allocations: HashMap::new(),
+            sub_loc_allocations: HashMap::new(),
+            ssa_val_allocations: func.ssa_vals.keys().map(|_| BTreeMap::new()).collect(),
         }
     }
     pub fn global_state(&self) -> &'a Rc<GlobalState> {
@@ -34,23 +36,30 @@ impl<'a> AllocatorState<'a> {
     pub fn live_ranges(&self) -> &[LiveRange] {
         &self.live_ranges
     }
-    pub fn allocations(
+    pub fn sub_loc_allocations(
         &self,
         sub_loc: SubLoc,
     ) -> &BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx> {
         const EMPTY: &'static BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx> = &BTreeMap::new();
-        self.allocations.get(&sub_loc).unwrap_or(EMPTY)
+        self.sub_loc_allocations.get(&sub_loc).unwrap_or(EMPTY)
+    }
+    pub fn ssa_val_allocations(
+        &self,
+        ssa_val_idx: SSAValIdx,
+    ) -> &BTreeMap<ProgRange<Lexicographic>, LiveRangeIdx> {
+        &self.ssa_val_allocations[ssa_val_idx.get()]
     }
     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) {
+                if let Some(allocations) = self.sub_loc_allocations.get_mut(&sub_loc) {
                     allocations.remove(&live_range.range);
                 }
             }
         }
+        self.ssa_val_allocations[live_range.ssa_val.get()].remove(&live_range.range);
         old_allocation
     }
     pub fn add_allocations_for(
@@ -63,12 +72,13 @@ impl<'a> AllocatorState<'a> {
         live_range.allocation = new_allocation;
         if let Some(new_allocation) = new_allocation {
             for sub_loc in new_allocation.sub_locs() {
-                self.allocations
+                self.sub_loc_allocations
                     .entry(sub_loc)
                     .or_default()
                     .insert(live_range.range, live_range_idx);
             }
         }
+        self.ssa_val_allocations[live_range.ssa_val.get()].insert(live_range.range, live_range_idx);
     }
     pub fn replace_allocation(
         &mut self,
-- 
2.30.2