Improve queue used by color_interference.

binary-banter · Dec 7, 2023 · 379735a · 379735a
1 parent 01e6a7a
commit 379735a
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 17 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/compiler/Cargo.toml b/compiler/Cargo.toml
@@ -18,6 +18,7 @@ derive_more = "0.99.17"
 functor_derive = "0.2.3"
 indenter = "0.3.3"
 once_cell = "1.18.0"
+binary-heap-plus = "0.5.0"
 
 [build-dependencies]
 lalrpop = "0.20.0"

diff --git a/compiler/src/passes/assign/color_interference.rs b/compiler/src/passes/assign/color_interference.rs
@@ -1,15 +1,21 @@
 use crate::passes::assign::{Arg, InterferenceGraph, LArg};
 use crate::passes::select::{Reg, CALLEE_SAVED_NO_STACK};
 use crate::utils::gen_sym::UniqueSym;
-use itertools::Itertools;
+use binary_heap_plus::BinaryHeap;
+use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 
 impl<'p> InterferenceGraph<'p> {
     #[must_use]
     pub fn color(self) -> (HashMap<UniqueSym<'p>, Arg>, usize) {
         let graph = self.0;
-        let mut queue = Vec::new();
-        let mut node_map = HashMap::<LArg, isize>::new();
+        let node_map = RefCell::new(HashMap::<LArg, isize>::new());
+        let mut queue = BinaryHeap::new_by_key(|node| {
+            graph
+                .neighbors(*node)
+                .filter(|nb| node_map.borrow().contains_key(nb))
+                .count()
+        });
 
         for node in graph.nodes() {
             match node {
@@ -35,33 +41,27 @@ impl<'p> InterferenceGraph<'p> {
                         Reg::R11 => -4,
                         Reg::R15 => -5,
                     };
-                    node_map.insert(node, node_weight);
+                    node_map.borrow_mut().insert(node, node_weight);
                 }
             }
         }
 
         while let Some(node) = queue.pop() {
             let used_colors = graph
                 .neighbors(node)
-                .filter_map(|nb| node_map.get(&nb))
+                .filter(|nb| node_map.borrow().contains_key(nb))
+                .map(|nb| node_map.borrow()[&nb])
                 .collect::<HashSet<_>>();
 
             let chosen_color = (0..)
                 .find(|i| !used_colors.contains(i))
                 .expect("there are infinite numbers, lol");
 
-            node_map.insert(node, chosen_color);
-
-            queue.sort_by_key(|node| {
-                graph
-                    .neighbors(*node)
-                    .filter_map(|nb| node_map.get(&nb))
-                    .unique()
-                    .count()
-            });
+            node_map.borrow_mut().insert(node, chosen_color);
         }
 
         let used_vars = node_map
+            .borrow()
             .values()
             .filter(|&&n| n >= 10)
             .map(|&n| n - 10)
@@ -71,6 +71,7 @@ impl<'p> InterferenceGraph<'p> {
         let stack_space = (8 * used_vars).div_ceil(16) * 16;
 
         let colors = node_map
+            .take()
             .into_iter()
             .filter_map(|(node, color)| match node {
                 LArg::Var { sym } => Some((sym, arg_from_color(color))),

diff --git a/compiler/src/passes/parse/tests.rs b/compiler/src/passes/parse/tests.rs
@@ -3,7 +3,7 @@ use crate::utils::split_test::split_test;
 use derive_name::VariantName;
 use test_each_file::test_each_file;
 
-fn parse([test]: [&str; 1]) {
+fn parsing([test]: [&str; 1]) {
     let (_, _, _, expected_error) = split_test(test);
 
     let result = parse(test);
@@ -22,5 +22,5 @@ fn parse([test]: [&str; 1]) {
     }
 }
 
-test_each_file! { for ["sp"] in "./programs/good" as parse_succeed => parse }
-test_each_file! { for ["sp"] in "./programs/fail/parse" as parse_fail => parse }
+test_each_file! { for ["sp"] in "./programs/good" as parse_succeed => parsing }
+test_each_file! { for ["sp"] in "./programs/fail/parse" as parse_fail => parsing }