Propagate register assignments.

dm/src/main.rs

@@ -61,7 +61,7 @@ fn main() {
         for declaration in compilation_unit.declarations() {
             if let ModuleLevelDeclaration::Function(function) = declaration {
                 let mut ir_function = function.to_ir(&symbol_table);
-                let register_assignments = ir_function.register_assignments(args.register_count);
+                let register_assignments = ir_function.assign_registers(args.register_count);
                 println!("{}", ir_function);
                 println!("{:?}", register_assignments);
             }

dmc-lib/src/ir/ir_add.rs

@@ -1,7 +1,7 @@
 use crate::ir::ir_expression::IrExpression;
 use crate::ir::ir_variable::IrVrVariableDescriptor;
 use crate::ir::register_allocation::VrUser;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 
 pub struct IrAdd {
@@ -43,4 +43,12 @@ impl VrUser for IrAdd {
         self.left.propagate_spills(spills);
         self.right.propagate_spills(spills);
     }
+
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    ) {
+        self.left.propagate_register_assignments(assignments);
+        self.right.propagate_register_assignments(assignments);
+    }
 }

dmc-lib/src/ir/ir_assign.rs

@@ -2,7 +2,7 @@ use crate::ir::ir_operation::IrOperation;
 use crate::ir::ir_variable::{IrVariable, IrVariableDescriptor, IrVrVariableDescriptor};
 use crate::ir::register_allocation::VrUser;
 use std::cell::RefCell;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
 
@@ -50,6 +50,21 @@ impl VrUser for IrAssign {
             }
         }
     }
+
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    ) {
+        let mut borrowed_destination = self.destination.borrow_mut();
+        if let IrVariableDescriptor::VirtualRegister(vr_variable) =
+            borrowed_destination.descriptor_mut()
+        {
+            if assignments.contains_key(vr_variable) {
+                vr_variable.set_assigned_register(assignments[vr_variable]);
+            }
+        }
+        self.initializer.propagate_register_assignments(assignments);
+    }
 }
 
 impl Display for IrAssign {

dmc-lib/src/ir/ir_block.rs

@@ -1,6 +1,6 @@
 use crate::ir::ir_statement::IrStatement;
 use crate::ir::ir_variable::IrVrVariableDescriptor;
-use crate::ir::register_allocation::{HasVrUsers, VrUser, registers_and_spills};
+use crate::ir::register_allocation::{HasVrUsers, VrUser};
 use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
@@ -32,34 +32,19 @@ impl IrBlock {
     pub fn statements(&self) -> &[IrStatement] {
         &self.statements
     }
-
-    pub fn register_assignments(
-        &mut self,
-        register_count: usize,
-    ) -> HashMap<IrVrVariableDescriptor, usize> {
-        let mut spills: HashSet<IrVrVariableDescriptor> = HashSet::new();
-        loop {
-            let mut interference_graph = self.interference_graph();
-            let (registers, new_spills) =
-                registers_and_spills(&mut interference_graph, register_count);
-
-            if spills != new_spills {
-                spills = new_spills;
-                // mutate all IrVirtualRegisters to constituent statements
-                for statement in &mut self.statements {
-                    statement.propagate_spills(&spills);
-                }
-            } else {
-                return registers;
-            }
-        }
-    }
 }
 
 impl HasVrUsers for IrBlock {
     fn vr_users(&self) -> Vec<&dyn VrUser> {
         self.statements.iter().map(|s| s as &dyn VrUser).collect()
     }
+
+    fn vr_users_mut(&mut self) -> Vec<&mut dyn VrUser> {
+        self.statements
+            .iter_mut()
+            .map(|s| s as &mut dyn VrUser)
+            .collect()
+    }
 }
 
 impl VrUser for IrBlock {
@@ -79,6 +64,15 @@ impl VrUser for IrBlock {
             statement.propagate_spills(spills);
         }
     }
+
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    ) {
+        for statement in &mut self.statements {
+            statement.propagate_register_assignments(assignments);
+        }
+    }
 }
 
 impl Display for IrBlock {
@@ -141,7 +135,7 @@ mod tests {
 
         if let ModuleLevelDeclaration::Function(main) = main {
             let mut main_ir = main.to_ir(&symbol_table);
-            let register_assignments = main_ir.register_assignments(2);
+            let register_assignments = main_ir.assign_registers(2);
             assert_eq!(register_assignments.len(), 4);
         } else {
             unreachable!()

dmc-lib/src/ir/ir_call.rs

@@ -1,7 +1,7 @@
 use crate::ir::ir_expression::IrExpression;
 use crate::ir::ir_variable::IrVrVariableDescriptor;
 use crate::ir::register_allocation::VrUser;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
 
@@ -33,6 +33,15 @@ impl VrUser for IrCall {
             argument.propagate_spills(spills);
         }
     }
+
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    ) {
+        for argument in &mut self.arguments {
+            argument.propagate_register_assignments(assignments);
+        }
+    }
 }
 
 impl Display for IrCall {

dmc-lib/src/ir/ir_expression.rs

@@ -4,7 +4,7 @@ use crate::ir::ir_variable::{
 };
 use crate::ir::register_allocation::VrUser;
 use std::cell::RefCell;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
 
@@ -85,4 +85,31 @@ impl VrUser for IrExpression {
             }
         }
     }
+
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    ) {
+        match self {
+            IrExpression::Parameter(_) => {
+                // no-op
+            }
+            IrExpression::Variable(ir_variable) => {
+                let mut borrowed = ir_variable.borrow_mut();
+                if let IrVariableDescriptor::VirtualRegister(vr_variable) =
+                    borrowed.descriptor_mut()
+                {
+                    if assignments.contains_key(vr_variable) {
+                        vr_variable.set_assigned_register(assignments[vr_variable]);
+                    }
+                }
+            }
+            IrExpression::Int(_) => {
+                // no-op
+            }
+            IrExpression::String(_) => {
+                // no-op
+            }
+        }
+    }
 }

dmc-lib/src/ir/ir_function.rs

@@ -1,6 +1,7 @@
 use crate::ir::ir_block::IrBlock;
 use crate::ir::ir_parameter::IrParameter;
 use crate::ir::ir_variable::IrVrVariableDescriptor;
+use crate::ir::register_allocation::{HasVrUsers, VrUser};
 use crate::type_info::TypeInfo;
 use std::cell::RefCell;
 use std::collections::HashMap;
@@ -29,11 +30,12 @@ impl IrFunction {
         }
     }
 
-    pub fn register_assignments(
+    /// Returns the register assignments for debugging purposes
+    pub fn assign_registers(
         &mut self,
         register_count: usize,
     ) -> HashMap<IrVrVariableDescriptor, usize> {
-        self.entry.borrow_mut().register_assignments(register_count)
+        self.entry.borrow_mut().assign_registers(register_count)
     }
 }
 

dmc-lib/src/ir/ir_operation.rs

@@ -3,7 +3,7 @@ use crate::ir::ir_call::IrCall;
 use crate::ir::ir_expression::IrExpression;
 use crate::ir::ir_variable::IrVrVariableDescriptor;
 use crate::ir::register_allocation::VrUser;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 
 pub enum IrOperation {
@@ -58,4 +58,21 @@ impl VrUser for IrOperation {
             }
         }
     }
+
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    ) {
+        match self {
+            IrOperation::Load(ir_expression) => {
+                ir_expression.propagate_register_assignments(assignments);
+            }
+            IrOperation::Add(ir_add) => {
+                ir_add.propagate_register_assignments(assignments);
+            }
+            IrOperation::Call(ir_call) => {
+                ir_call.propagate_register_assignments(assignments);
+            }
+        }
+    }
 }

dmc-lib/src/ir/ir_return.rs

@@ -1,7 +1,7 @@
 use crate::ir::ir_expression::IrExpression;
 use crate::ir::ir_variable::IrVrVariableDescriptor;
 use crate::ir::register_allocation::VrUser;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 
 pub struct IrReturn {
@@ -32,6 +32,15 @@ impl VrUser for IrReturn {
             ir_expression.propagate_spills(spills);
         }
     }
+
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    ) {
+        if let Some(ir_expression) = self.value.as_mut() {
+            ir_expression.propagate_register_assignments(assignments);
+        }
+    }
 }
 
 impl Display for IrReturn {

dmc-lib/src/ir/ir_statement.rs

@@ -3,7 +3,7 @@ use crate::ir::ir_call::IrCall;
 use crate::ir::ir_return::IrReturn;
 use crate::ir::ir_variable::IrVrVariableDescriptor;
 use crate::ir::register_allocation::VrUser;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
 
 pub enum IrStatement {
@@ -42,6 +42,23 @@ impl VrUser for IrStatement {
             }
         }
     }
+
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    ) {
+        match self {
+            IrStatement::Assign(ir_assign) => {
+                ir_assign.propagate_register_assignments(assignments);
+            }
+            IrStatement::Call(ir_call) => {
+                ir_call.propagate_register_assignments(assignments);
+            }
+            IrStatement::Return(ir_return) => {
+                ir_return.propagate_register_assignments(assignments);
+            }
+        }
+    }
 }
 
 impl Display for IrStatement {

dmc-lib/src/ir/ir_variable.rs

@@ -33,6 +33,10 @@ impl IrVariable {
         &self.descriptor
     }
 
+    pub fn descriptor_mut(&mut self) -> &mut IrVariableDescriptor {
+        &mut self.descriptor
+    }
+
     pub fn set_descriptor(&mut self, descriptor: IrVariableDescriptor) {
         self.descriptor = descriptor;
     }
@@ -94,6 +98,10 @@ impl IrVrVariableDescriptor {
     pub fn block_id(&self) -> usize {
         self.block_id
     }
+
+    pub fn set_assigned_register(&mut self, register: usize) {
+        self.assigned_register = Some(register);
+    }
 }
 
 impl Display for IrVrVariableDescriptor {

dmc-lib/src/ir/register_allocation.rs

@@ -6,6 +6,7 @@ pub type LivenessMap = HashMap<usize, HashSet<IrVrVariableDescriptor>>;
 
 pub trait HasVrUsers {
     fn vr_users(&self) -> Vec<&dyn VrUser>;
+    fn vr_users_mut(&mut self) -> Vec<&mut dyn VrUser>;
 
     fn live_in_live_out(&self) -> (LivenessMap, LivenessMap) {
         let mut live_in: LivenessMap = HashMap::new();
@@ -99,12 +100,42 @@ pub trait HasVrUsers {
 
         graph
     }
+
+    fn assign_registers(
+        &mut self,
+        register_count: usize,
+    ) -> HashMap<IrVrVariableDescriptor, usize> {
+        let mut spills: HashSet<IrVrVariableDescriptor> = HashSet::new();
+        loop {
+            let mut interference_graph = self.interference_graph();
+            let (registers, new_spills) =
+                registers_and_spills(&mut interference_graph, register_count);
+
+            if spills != new_spills {
+                spills = new_spills;
+                // propagate all spills, since those won't be used for the next interference graph
+                for vr_user in &mut self.vr_users_mut() {
+                    vr_user.propagate_spills(&spills);
+                }
+            } else {
+                // we've calculated final assignments, so propagate them
+                for vr_user in self.vr_users_mut() {
+                    vr_user.propagate_register_assignments(&registers);
+                }
+                return registers;
+            }
+        }
+    }
 }
 
 pub trait VrUser {
     fn vr_definitions(&self) -> HashSet<IrVrVariableDescriptor>;
     fn vr_uses(&self) -> HashSet<IrVrVariableDescriptor>;
     fn propagate_spills(&mut self, spills: &HashSet<IrVrVariableDescriptor>);
+    fn propagate_register_assignments(
+        &mut self,
+        assignments: &HashMap<IrVrVariableDescriptor, usize>,
+    );
 }
 
 #[derive(Debug)]