Fix register allocation algorithm.

dm/src/main.rs

@@ -25,6 +25,9 @@ struct Cli {
 
     #[arg(long)]
     show_ir: bool,
+
+    #[arg(long, default_value = "8")]
+    register_count: usize,
 }
 
 fn main() {
@@ -58,8 +61,9 @@ fn main() {
         for declaration in compilation_unit.declarations() {
             if let ModuleLevelDeclaration::Function(function) = declaration {
                 let mut ir_function = function.to_ir(&symbol_table);
-                ir_function.assign_registers();
+                let register_assignments = ir_function.register_assignments(args.register_count);
-                println!("{}", ir_function)
+                println!("{}", ir_function);
+                println!("{:?}", register_assignments);
             }
         }
     }

dmc-lib/src/ir/ir_add.rs

@@ -23,6 +23,11 @@ impl IrAdd {
         set.extend(self.right.vr_uses());
         set
     }
+
+    pub fn propagate_spills(&mut self, spills: &HashSet<Rc<IrVirtualRegisterVariable>>) {
+        self.left.propagate_spills(spills);
+        self.right.propagate_spills(spills);
+    }
 }
 
 impl Display for IrAdd {

dmc-lib/src/ir/ir_assign.rs

@@ -34,8 +34,10 @@ impl IrAssign {
     }
 
     pub fn propagate_spills(&mut self, spills: &HashSet<Rc<IrVirtualRegisterVariable>>) {
+        self.initializer.propagate_spills(spills);
         if let IrVariable::VirtualRegister(vr_variable) = self.destination.deref() {
             if spills.contains(vr_variable) {
+                println!("changing vr to stack: {}", vr_variable.name());
                 self.destination = Box::new(IrVariable::Stack(IrStackVariable::new(
                     vr_variable.name(),
                     vr_variable.type_info().clone(),

dmc-lib/src/ir/ir_block.rs

@@ -14,6 +14,8 @@ pub struct IrBlock {
 }
 
 type LivenessMapByStatement = HashMap<usize, HashSet<Rc<IrVirtualRegisterVariable>>>;
+type InterferenceGraph =
+    HashMap<Rc<IrVirtualRegisterVariable>, HashSet<Rc<IrVirtualRegisterVariable>>>;
 
 impl IrBlock {
     pub fn new(id: usize, debug_label: &str, statements: Vec<IrStatement>) -> Self {
@@ -98,19 +100,14 @@ impl IrBlock {
         (live_in, live_out)
     }
 
-    fn interference_graph(
+    fn interference_graph(&self) -> InterferenceGraph {
-        &self,
-    ) -> HashMap<Rc<IrVirtualRegisterVariable>, HashSet<Rc<IrVirtualRegisterVariable>>> {
         let mut all_vr_variables: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
         for statement in &self.statements {
             all_vr_variables.extend(statement.vr_definitions());
             all_vr_variables.extend(statement.vr_uses());
         }
 
-        let mut graph: HashMap<
+        let mut graph: InterferenceGraph = HashMap::new();
-            Rc<IrVirtualRegisterVariable>,
-            HashSet<Rc<IrVirtualRegisterVariable>>,
-        > = HashMap::new();
         for variable in all_vr_variables {
             graph.insert(variable, HashSet::new());
         }
@@ -139,10 +136,16 @@ impl IrBlock {
         graph
     }
 
-    pub fn assign_registers(&mut self) {
+    pub fn register_assignments(
+        &mut self,
+        register_count: usize,
+    ) -> HashMap<Rc<IrVirtualRegisterVariable>, usize> {
         let mut spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
         loop {
-            let (registers, new_spills) = register_assignment::assign_registers(self);
+            let mut interference_graph = self.interference_graph();
+            let (registers, new_spills) =
+                register_assignment::registers_and_spills(&mut interference_graph, register_count);
+
             if spills != new_spills {
                 spills = new_spills;
                 // mutate all IrVirtualRegisters to constituent statements
@@ -150,8 +153,7 @@ impl IrBlock {
                     statement.propagate_spills(&spills);
                 }
             } else {
-                println!("{:?}", registers);
+                return registers;
-                break;
             }
         }
     }
@@ -185,6 +187,46 @@ impl Display for IrBlock {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use crate::ast::module_level_declaration::ModuleLevelDeclaration;
+    use crate::parser::parse_compilation_unit;
+    use crate::symbol_table::SymbolTable;
+
+    #[test]
+    fn overlapping_assignments_bug_when_k_2() {
+        let mut compilation_unit = parse_compilation_unit(
+            "
+            fn main()
+                let a = 1
+                let b = 2
+                let c = 3
+                let x = a + b + c
+            end
+        ",
+        )
+        .unwrap();
+        let mut symbol_table = SymbolTable::new();
+        compilation_unit.gather_declared_names(&mut symbol_table);
+        compilation_unit.check_name_usages(&mut symbol_table);
+        compilation_unit.type_check(&mut symbol_table);
+
+        let main = compilation_unit
+            .declarations()
+            .iter()
+            .find(|d| matches!(d, ModuleLevelDeclaration::Function(_)))
+            .unwrap();
+
+        if let ModuleLevelDeclaration::Function(main) = main {
+            let mut main_ir = main.to_ir(&symbol_table);
+            let register_assignments = main_ir.register_assignments(2);
+            assert_eq!(register_assignments.len(), 4);
+        } else {
+            unreachable!()
+        }
+    }
+}
+
 mod register_assignment {
     use super::*;
 
@@ -195,83 +237,158 @@ mod register_assignment {
         color: bool,
     }
 
-    pub fn assign_registers(
+    pub fn registers_and_spills(
-        block: &IrBlock,
+        interference_graph: &mut InterferenceGraph,
+        k: usize,
     ) -> (
         HashMap<Rc<IrVirtualRegisterVariable>, usize>,
         HashSet<Rc<IrVirtualRegisterVariable>>,
     ) {
-        let k = 8;
-        let mut spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
-
-        loop {
-            // 1. get interference graph
-            let mut interference_graph = block.interference_graph();
         let mut work_stack: Vec<WorkItem> = vec![];
 
-            loop {
+        while !interference_graph.is_empty() {
-                // 2. coloring by simplification
+            let next = next_work_item(interference_graph, k);
-                //    try to find a node (virtual register) with less than k outgoing edges,
+            work_stack.push(next);
-                //    and mark as color
-                //    if not, pick any, and mark as spill for step 3
-                let register_lt_k = interference_graph.iter().find_map(|(vr, neighbors)| {
-                    if neighbors.len() < k { Some(vr) } else { None }
-                });
-                if let Some(vr) = register_lt_k {
-                    let vr = vr.clone();
-                    // remove both outgoing and incoming edges; save either set for WorkItem
-                    // first, outgoing:
-                    let outgoing_edges = interference_graph.remove(&vr).unwrap();
-
-                    // second, incoming
-                    interference_graph.iter_mut().for_each(|(_, neighbors)| {
-                        neighbors.remove(&vr);
-                    });
-
-                    // push to work stack
-                    work_stack.push(WorkItem {
-                        vr,
-                        edges: outgoing_edges,
-                        color: true,
-                    })
-                } else {
-                    // pick any
-                    let vr = interference_graph.iter().last().unwrap().0.clone();
-
-                    // first, outgoing
-                    let outgoing_edges = interference_graph.remove(&vr).unwrap();
-
-                    // second, incoming
-                    interference_graph.iter_mut().for_each(|(_, neighbors)| {
-                        neighbors.remove(&vr);
-                    });
-
-                    work_stack.push(WorkItem {
-                        vr,
-                        edges: outgoing_edges,
-                        color: false, // spill
-                    });
-                }
-
-                if interference_graph.is_empty() {
-                    break;
-                }
         }
 
         // 3. assign colors to registers
-            let mut rebuilt_graph: HashMap<
+        let mut rebuilt_graph: InterferenceGraph = HashMap::new();
-                Rc<IrVirtualRegisterVariable>,
-                HashSet<Rc<IrVirtualRegisterVariable>>,
-            > = HashMap::new();
         let mut register_assignments: HashMap<Rc<IrVirtualRegisterVariable>, usize> =
             HashMap::new();
-            let mut new_spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
+        let mut spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
 
         while let Some(work_item) = work_stack.pop() {
             if work_item.color {
                 assign_register(&work_item, &mut rebuilt_graph, k, &mut register_assignments);
+            } else if can_optimistically_color(&work_item, &mut register_assignments, k) {
+                assign_register(&work_item, &mut rebuilt_graph, k, &mut register_assignments);
             } else {
-                    // first, see if we can optimistically color
+                // spill
+                spills.insert(work_item.vr.clone());
+            }
+        }
+
+        (register_assignments, spills)
+    }
+
+    fn assign_register(
+        work_item: &WorkItem,
+        graph: &mut InterferenceGraph,
+        k: usize,
+        register_assignments: &mut HashMap<Rc<IrVirtualRegisterVariable>, usize>,
+    ) {
+        rebuild_vr_and_edges(graph, work_item);
+
+        // find a register which is not yet shared by all outgoing edges' vertices
+        'outer: for i in 0..k {
+            for edge in graph.get_mut(&work_item.vr).unwrap().iter() {
+                if register_assignments.contains_key(edge) {
+                    let assignment = register_assignments.get(edge).unwrap();
+                    if *assignment == i {
+                        continue 'outer;
+                    }
+                }
+            }
+            register_assignments.insert(work_item.vr.clone(), i);
+            break;
+        }
+    }
+
+    fn find_vr_lt_k(
+        interference_graph: &InterferenceGraph,
+        k: usize,
+    ) -> Option<&Rc<IrVirtualRegisterVariable>> {
+        interference_graph.iter().find_map(
+            |(vr, neighbors)| {
+                if neighbors.len() < k { Some(vr) } else { None }
+            },
+        )
+    }
+
+    fn remove_vr_and_edges(
+        interference_graph: &mut InterferenceGraph,
+        vr: &Rc<IrVirtualRegisterVariable>,
+    ) -> HashSet<Rc<IrVirtualRegisterVariable>> {
+        // first, outgoing
+        let outgoing_edges = interference_graph.remove(vr).unwrap();
+
+        // second, incoming
+        for neighbor in &outgoing_edges {
+            let neighbor_edges = interference_graph.get_mut(neighbor).unwrap();
+            neighbor_edges.remove(vr);
+        }
+
+        outgoing_edges
+    }
+
+    fn next_work_item(interference_graph: &mut InterferenceGraph, k: usize) -> WorkItem {
+        // try to find a node (virtual register) with less than k outgoing edges, and mark as color
+        // for step 3.
+        // if not, pick any, and mark as spill for step 3.
+        let register_lt_k = find_vr_lt_k(interference_graph, k);
+        if let Some(vr) = register_lt_k {
+            let vr = vr.clone();
+
+            // remove edges; save outgoing to work_item
+            let edges = remove_vr_and_edges(interference_graph, &vr);
+
+            // push to work stack
+            WorkItem {
+                vr,
+                edges,
+                color: true,
+            }
+        } else {
+            // pick any
+            let vr = interference_graph.iter().last().unwrap().0.clone();
+
+            // remove edges
+            let edges = remove_vr_and_edges(interference_graph, &vr);
+
+            WorkItem {
+                vr,
+                edges,
+                color: false, // spill
+            }
+        }
+    }
+
+    fn rebuild_vr_and_edges(graph: &mut InterferenceGraph, work_item: &WorkItem) {
+        // init the vertex
+        graph.insert(work_item.vr.clone(), HashSet::new());
+
+        // outgoing
+        for neighbor in &work_item.edges {
+            // check if neighbor exists in the graph first; if it was marked spill earlier and could
+            // not optimistically color, it won't be in the graph
+            if graph.contains_key(neighbor) {
+                // get outgoing set and insert neighbor
+                graph
+                    .get_mut(&work_item.vr)
+                    .unwrap()
+                    .insert(neighbor.clone());
+            }
+        }
+
+        // incoming
+        for neighbor in &work_item.edges {
+            // like above, neighbor may not have been added because of failure to optimistically
+            // color
+            if graph.contains_key(neighbor) {
+                graph
+                    .get_mut(neighbor)
+                    .unwrap()
+                    .insert(work_item.vr.clone());
+            }
+        }
+    }
+
+    fn can_optimistically_color(
+        work_item: &WorkItem,
+        register_assignments: &HashMap<Rc<IrVirtualRegisterVariable>, usize>,
+        k: usize,
+    ) -> bool {
+        // see if we can optimistically color
         // find how many assignments have been made for the outgoing edges
         // if it's less than k, we can do it
         let mut number_of_assigned_edges = 0;
@@ -280,72 +397,135 @@ mod register_assignment {
                 number_of_assigned_edges += 1;
             }
         }
-
+        number_of_assigned_edges < k
-                    if number_of_assigned_edges < k {
-                        // optimistically color
-                        assign_register(
-                            &work_item,
-                            &mut rebuilt_graph,
-                            k,
-                            &mut register_assignments,
-                        );
-                    } else {
-                        // spill
-                        new_spills.insert(work_item.vr.clone());
     }
+
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+        use crate::type_info::TypeInfo;
+
+        fn line_graph() -> InterferenceGraph {
+            let mut graph: InterferenceGraph = HashMap::new();
+            let v0 = Rc::new(IrVirtualRegisterVariable::new("v0", TypeInfo::Integer));
+            let v1 = Rc::new(IrVirtualRegisterVariable::new("v1", TypeInfo::Integer));
+            let v2 = Rc::new(IrVirtualRegisterVariable::new("v2", TypeInfo::Integer));
+
+            // v1 -- v0 -- v2
+            graph.insert(v0.clone(), HashSet::from([v1.clone(), v2.clone()]));
+            graph.insert(v1.clone(), HashSet::from([v0.clone()]));
+            graph.insert(v2.clone(), HashSet::from([v0.clone()]));
+            graph
+        }
+
+        fn triangle_graph() -> InterferenceGraph {
+            let mut graph: InterferenceGraph = HashMap::new();
+            let v0 = Rc::new(IrVirtualRegisterVariable::new("v0", TypeInfo::Integer));
+            let v1 = Rc::new(IrVirtualRegisterVariable::new("v1", TypeInfo::Integer));
+            let v2 = Rc::new(IrVirtualRegisterVariable::new("v2", TypeInfo::Integer));
+
+            // triangle: each has two edges
+            // v0
+            // |  \
+            // v1--v2
+            graph.insert(v0.clone(), HashSet::from([v1.clone(), v2.clone()]));
+            graph.insert(v1.clone(), HashSet::from([v0.clone(), v2.clone()]));
+            graph.insert(v2.clone(), HashSet::from([v0.clone(), v1.clone()]));
+            graph
+        }
+
+        fn get_vrs(graph: &InterferenceGraph) -> Vec<Rc<IrVirtualRegisterVariable>> {
+            let v0 = graph.keys().find(|k| k.name() == "v0").unwrap();
+            let v1 = graph.keys().find(|k| k.name() == "v1").unwrap();
+            let v2 = graph.keys().find(|k| k.name() == "v2").unwrap();
+            vec![v0.clone(), v1.clone(), v2.clone()]
+        }
+
+        #[test]
+        fn find_vr_lt_k_when_k_2() {
+            let graph = line_graph();
+            let found = find_vr_lt_k(&graph, 2);
+            assert!(found.is_some());
+            assert!(found.unwrap().name() == "v1" || found.unwrap().name() == "v2");
+        }
+
+        #[test]
+        fn find_vr_lt_k_when_k_1() {
+            let graph = line_graph();
+            let found = find_vr_lt_k(&graph, 1);
+            assert!(found.is_none());
+        }
+
+        #[test]
+        fn remove_edges_v0() {
+            let mut graph = line_graph();
+            let vrs = get_vrs(&graph);
+
+            let v0_outgoing = remove_vr_and_edges(&mut graph, &vrs[0]);
+            assert!(v0_outgoing.contains(&vrs[1]));
+            assert!(v0_outgoing.contains(&vrs[2]));
+
+            // check that incoming edges were removed
+            let v1_outgoing = graph.get(&vrs[1]).unwrap();
+            assert!(v1_outgoing.is_empty());
+            let v2_outgoing = graph.get(&vrs[2]).unwrap();
+            assert!(v2_outgoing.is_empty());
+        }
+
+        fn triangle_work_stack_k_2() -> Vec<WorkItem> {
+            let k = 2;
+            let mut graph = triangle_graph();
+
+            let mut work_stack = vec![];
+
+            // run three times, once for each register
+            work_stack.push(next_work_item(&mut graph, k));
+            work_stack.push(next_work_item(&mut graph, k));
+            work_stack.push(next_work_item(&mut graph, k));
+
+            work_stack
+        }
+
+        #[test]
+        fn next_work_item_k_2() {
+            let work_stack = triangle_work_stack_k_2();
+
+            // the actual edges may be different, depending on the underlying order in the sets
+            // (HashSet seems to use randomness in order)
+            // however, the bottommost item must be a spill, and the edge counts must be (from the
+            // bottom of the stack) 2-1-0
+            assert!(!work_stack[0].color);
+            assert_eq!(work_stack[0].edges.len(), 2);
+            assert_eq!(work_stack[1].edges.len(), 1);
+            assert_eq!(work_stack[2].edges.len(), 0);
+        }
+
+        #[test]
+        fn rebuild_graph_triangle_k_2() {
+            let mut work_stack = triangle_work_stack_k_2();
+            let mut rebuilt_graph: InterferenceGraph = HashMap::new();
+
+            // it should be possible to rebuild the graph from the stack, without yet worrying
+            // about spilling/etc.
+            while let Some(work_item) = work_stack.pop() {
+                rebuild_vr_and_edges(&mut rebuilt_graph, &work_item);
+            }
+
+            // we should have a triangle graph again
+            let vrs = get_vrs(&rebuilt_graph);
+            for vr in &vrs {
+                assert!(rebuilt_graph.contains_key(vr));
+                assert_eq!(rebuilt_graph.get(vr).unwrap().len(), 2);
             }
         }
 
-            if spills.eq(&new_spills) {
+        #[test]
-                return (register_assignments, spills);
+        fn registers_and_spills_triangle_k_2() {
-            } else {
+            let mut graph = triangle_graph();
-                spills = new_spills;
+            let (registers, spills) = registers_and_spills(&mut graph, 2);
-            }
+            // there should be one spill when k is 2
-        }
+            assert_eq!(registers.len(), 2);
-    }
+            assert_eq!(spills.len(), 1);
-
-    fn assign_register(
-        work_item: &WorkItem,
-        rebuilt_graph: &mut HashMap<
-            Rc<IrVirtualRegisterVariable>,
-            HashSet<Rc<IrVirtualRegisterVariable>>,
-        >,
-        k: usize,
-        register_assignments: &mut HashMap<Rc<IrVirtualRegisterVariable>, usize>,
-    ) {
-        let this_vertex_vr = &work_item.vr;
-        // init the vertex
-        rebuilt_graph.insert(this_vertex_vr.clone(), HashSet::new());
-
-        // add edges, both outgoing and incoming
-        let neighbors = rebuilt_graph.get_mut(this_vertex_vr).unwrap();
-        // outgoing
-        for edge in &work_item.edges {
-            neighbors.insert(edge.clone());
-        }
-        // incoming
-        for neighbor in neighbors.clone() {
-            if rebuilt_graph.contains_key(&neighbor) {
-                rebuilt_graph
-                    .get_mut(&neighbor)
-                    .unwrap()
-                    .insert(this_vertex_vr.clone());
-            }
-        }
-
-        // find a register which is not yet shared by all outgoing edges' vertices
-        // I think the bug is somewhere here
-        'outer: for i in 0..k {
-            for edge in rebuilt_graph.get_mut(this_vertex_vr).unwrap().iter() {
-                if register_assignments.contains_key(edge) {
-                    let assignment = register_assignments.get(edge).unwrap();
-                    if *assignment == i {
-                        continue 'outer;
-                    }
-                }
-            }
-            register_assignments.insert(this_vertex_vr.clone(), i);
-            break;
         }
     }
 }

dmc-lib/src/ir/ir_function.rs

@@ -1,7 +1,9 @@
 use crate::ir::ir_block::IrBlock;
 use crate::ir::ir_parameter::IrParameter;
+use crate::ir::ir_variable::IrVirtualRegisterVariable;
 use crate::type_info::TypeInfo;
 use std::cell::RefCell;
+use std::collections::HashMap;
 use std::fmt::Display;
 use std::rc::Rc;
 
@@ -27,8 +29,11 @@ impl IrFunction {
         }
     }
 
-    pub fn assign_registers(&mut self) {
+    pub fn register_assignments(
-        self.entry.borrow_mut().assign_registers();
+        &mut self,
+        register_count: usize,
+    ) -> HashMap<Rc<IrVirtualRegisterVariable>, usize> {
+        self.entry.borrow_mut().register_assignments(register_count)
     }
 }
 

dmc-lib/src/ir/ir_operation.rs

@@ -1,7 +1,7 @@
 use crate::ir::ir_add::IrAdd;
 use crate::ir::ir_call::IrCall;
 use crate::ir::ir_expression::IrExpression;
-use crate::ir::ir_variable::{IrVariable, IrVirtualRegisterVariable};
+use crate::ir::ir_variable::IrVirtualRegisterVariable;
 use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
@@ -36,4 +36,12 @@ impl IrOperation {
             IrOperation::Call(ir_call) => ir_call.vr_uses(),
         }
     }
+
+    pub fn propagate_spills(&mut self, spills: &HashSet<Rc<IrVirtualRegisterVariable>>) {
+        match self {
+            IrOperation::Load(ir_expression) => ir_expression.propagate_spills(spills),
+            IrOperation::Add(ir_add) => ir_add.propagate_spills(spills),
+            IrOperation::Call(ir_call) => ir_call.propagate_spills(spills),
+        }
+    }
 }