From 379735affb0f32be1a70ff45dfa37de7fcaa045d Mon Sep 17 00:00:00 2001
From: Vlamonster <jellemartijndijkstra@gmail.com>
Date: Thu, 7 Dec 2023 17:17:36 +0100
Subject: [PATCH] Improve queue used by `color_interference`.

---
 Cargo.lock                                    | 16 ++++++++++
 compiler/Cargo.toml                           |  1 +
 .../src/passes/assign/color_interference.rs   | 29 ++++++++++---------
 compiler/src/passes/parse/tests.rs            |  6 ++--
 4 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 6affbeb..a6f4d26 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -391,6 +391,15 @@ dependencies = [
  "walkdir",
 ]
 
+[[package]]
+name = "binary-heap-plus"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e4551d8382e911ecc0d0f0ffb602777988669be09447d536ff4388d1def11296"
+dependencies = [
+ "compare",
+]
+
 [[package]]
 name = "bit-set"
 version = "0.5.3"
@@ -568,10 +577,17 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7"
 
+[[package]]
+name = "compare"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "120133d4db2ec47efe2e26502ee984747630c67f51974fca0b6c1340cf2368d3"
+
 [[package]]
 name = "compiler"
 version = "0.1.0"
 dependencies = [
+ "binary-heap-plus",
  "bitflags 2.4.1",
  "clap",
  "derive-name",
diff --git a/compiler/Cargo.toml b/compiler/Cargo.toml
index be289d4..3549b3e 100644
--- 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
index 235a031..6dc0fb3 100644
--- 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,7 +41,7 @@ impl<'p> InterferenceGraph<'p> {
                         Reg::R11 => -4,
                         Reg::R15 => -5,
                     };
-                    node_map.insert(node, node_weight);
+                    node_map.borrow_mut().insert(node, node_weight);
                 }
             }
         }
@@ -43,25 +49,19 @@ impl<'p> InterferenceGraph<'p> {
         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
index 0bd8e31..9977b7e 100644
--- 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 }