Implement register allocation algorithm, with a few bugs.

dm/src/main.rs

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

dmc-lib/src/ast/ast_to_ir_util.rs

@@ -4,9 +4,10 @@ use crate::ir::ir_assign::IrAssign;
 use crate::ir::ir_expression::IrExpression;
 use crate::ir::ir_operation::IrOperation;
 use crate::ir::ir_statement::IrStatement;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::{IrVariable, IrVirtualRegisterVariable};
 use crate::symbol_table::SymbolTable;
 use crate::type_info::TypeInfo;
+use std::rc::Rc;
 
 pub fn expression_to_ir_expression(
     expression: &Expression,
@@ -22,13 +23,13 @@ pub fn expression_to_ir_expression(
                     .add_statement(IrStatement::Call(ir_call));
                 None
             } else {
-                let t_var = IrVariable::new(&builder.new_t_var(), call.type_info());
+                let t_var = IrVirtualRegisterVariable::new(&builder.new_t_var(), call.type_info());
-                let assign = IrAssign::new(t_var, IrOperation::Call(ir_call));
+                let as_rc = Rc::new(t_var);
-                let destination = assign.destination().clone();
+                let assign = IrAssign::new(as_rc.clone(), IrOperation::Call(ir_call));
                 builder
                     .current_block_mut()
                     .add_statement(IrStatement::Assign(assign));
-                Some(IrExpression::Variable(destination))
+                Some(IrExpression::Variable(IrVariable::VirtualRegister(as_rc)))
             }
         }
         Expression::IntegerLiteral(integer_literal) => {
@@ -43,13 +44,16 @@ pub fn expression_to_ir_expression(
         }
         Expression::Additive(additive_expression) => {
             let ir_add = additive_expression.to_ir(builder, symbol_table);
-            let t_var = IrVariable::new(&builder.new_t_var(), additive_expression.type_info());
+            let t_var = IrVirtualRegisterVariable::new(
-            let assign = IrAssign::new(t_var, IrOperation::Add(ir_add));
+                &builder.new_t_var(),
-            let destination = assign.destination().clone();
+                additive_expression.type_info(),
+            );
+            let as_rc = Rc::new(t_var);
+            let assign = IrAssign::new(as_rc.clone(), IrOperation::Add(ir_add));
             builder
                 .current_block_mut()
                 .add_statement(IrStatement::Assign(assign));
-            Some(IrExpression::Variable(destination))
+            Some(IrExpression::Variable(IrVariable::VirtualRegister(as_rc)))
         }
     }
 }

dmc-lib/src/ast/function.rs

@@ -8,7 +8,6 @@ use crate::constants_table::ConstantsTable;
 use crate::diagnostic::Diagnostic;
 use crate::ir::ir_function::IrFunction;
 use crate::ir::ir_parameter::IrParameter;
-use crate::ir::ir_variable::IrVariable;
 use crate::source_range::SourceRange;
 use crate::symbol::FunctionSymbol;
 use crate::symbol_table::{SymbolInsertError, SymbolTable};

dmc-lib/src/ast/let_statement.rs

@@ -8,10 +8,11 @@ use crate::ir::ir_assign::IrAssign;
 use crate::ir::ir_expression::IrExpression;
 use crate::ir::ir_operation::IrOperation;
 use crate::ir::ir_statement::IrStatement;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::{IrVariable, IrVirtualRegisterVariable};
 use crate::source_range::SourceRange;
 use crate::symbol::{ExpressibleSymbol, VariableSymbol};
 use crate::symbol_table::{SymbolInsertError, SymbolTable};
+use std::rc::Rc;
 
 pub struct LetStatement {
     declared_name: String,
@@ -106,14 +107,11 @@ impl LetStatement {
 
         let destination_symbol =
             symbol_table.get_variable_symbol(self.scope_id.unwrap(), &self.declared_name);
-        let ir_assign = IrAssign::new(
+        let destination_vr_variable =
-            IrVariable::new(self.declared_name(), self.initializer().type_info()),
+            IrVirtualRegisterVariable::new(self.declared_name(), self.initializer().type_info());
-            init_operation,
+        let as_rc = Rc::new(destination_vr_variable);
-        );
+        let ir_assign = IrAssign::new(as_rc.clone(), init_operation);
-        let destination_ir_variable = ir_assign.destination();
+        destination_symbol.borrow_mut().set_vr_variable(as_rc);
-        destination_symbol
-            .borrow_mut()
-            .set_ir_variable(destination_ir_variable.clone());
 
         builder
             .current_block_mut()

dmc-lib/src/ir/ir_add.rs

@@ -1,5 +1,5 @@
 use crate::ir::ir_expression::IrExpression;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::IrVirtualRegisterVariable;
 use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
@@ -16,6 +16,13 @@ impl IrAdd {
             right: right.into(),
         }
     }
+
+    pub fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
+        let mut set = HashSet::new();
+        set.extend(self.left.vr_uses());
+        set.extend(self.right.vr_uses());
+        set
+    }
 }
 
 impl Display for IrAdd {
@@ -23,12 +30,3 @@ impl Display for IrAdd {
         write!(f, "{} + {}", self.left, self.right)
     }
 }
-
-impl IrAdd {
-    pub fn uses(&self) -> HashSet<Rc<IrVariable>> {
-        let mut set = HashSet::new();
-        set.extend(self.left.uses());
-        set.extend(self.right.uses());
-        set
-    }
-}

dmc-lib/src/ir/ir_assign.rs

@@ -1,32 +1,47 @@
 use crate::ir::ir_operation::IrOperation;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::{IrStackVariable, IrVariable, IrVirtualRegisterVariable};
 use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
+use std::ops::Deref;
 use std::rc::Rc;
 
 pub struct IrAssign {
-    destination: Rc<IrVariable>,
+    destination: Box<IrVariable>,
     initializer: Box<IrOperation>,
 }
 
 impl IrAssign {
-    pub fn new(destination: IrVariable, initializer: IrOperation) -> Self {
+    pub fn new(destination: Rc<IrVirtualRegisterVariable>, initializer: IrOperation) -> Self {
         Self {
-            destination: Rc::new(destination),
+            destination: Box::new(IrVariable::VirtualRegister(destination)),
             initializer: initializer.into(),
         }
     }
 
-    pub fn destination(&self) -> &Rc<IrVariable> {
+    pub fn vr_definitions(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
-        &self.destination
+        match self.destination.deref() {
+            IrVariable::VirtualRegister(vr_variable) => {
+                let mut set = HashSet::new();
+                set.insert(vr_variable.clone());
+                set
+            }
+            IrVariable::Stack(_) => HashSet::new(),
+        }
     }
 
-    pub fn initializer(&self) -> &IrOperation {
+    pub fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
-        &self.initializer
+        self.initializer.vr_uses()
     }
 
-    pub fn uses(&self) -> HashSet<Rc<IrVariable>> {
+    pub fn propagate_spills(&mut self, spills: &HashSet<Rc<IrVirtualRegisterVariable>>) {
-        self.initializer.uses()
+        if let IrVariable::VirtualRegister(vr_variable) = self.destination.deref() {
+            if spills.contains(vr_variable) {
+                self.destination = Box::new(IrVariable::Stack(IrStackVariable::new(
+                    vr_variable.name(),
+                    vr_variable.type_info().clone(),
+                )));
+            }
+        }
     }
 }
 

dmc-lib/src/ir/ir_block.rs

@@ -1,5 +1,5 @@
 use crate::ir::ir_statement::IrStatement;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::IrVirtualRegisterVariable;
 use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::fmt::{Display, Formatter};
@@ -13,6 +13,8 @@ pub struct IrBlock {
     statements: Vec<IrStatement>,
 }
 
+type LivenessMapByStatement = HashMap<usize, HashSet<Rc<IrVirtualRegisterVariable>>>;
+
 impl IrBlock {
     pub fn new(id: usize, debug_label: &str, statements: Vec<IrStatement>) -> Self {
         Self {
@@ -32,30 +34,25 @@ impl IrBlock {
         &self.statements
     }
 
-    fn definitions(&self) -> HashSet<Rc<IrVariable>> {
+    fn vr_definitions(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
         let mut set = HashSet::new();
         for statement in &self.statements {
-            set.extend(statement.definitions())
+            set.extend(statement.vr_definitions());
         }
         set
     }
 
-    fn uses(&self) -> HashSet<Rc<IrVariable>> {
+    fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
         let mut set = HashSet::new();
         for statement in &self.statements {
-            set.extend(statement.uses());
+            set.extend(statement.vr_uses());
         }
         set
     }
 
-    fn live_in_live_out(
+    fn live_in_live_out(&self) -> (LivenessMapByStatement, LivenessMapByStatement) {
-        &self,
+        let mut live_in: LivenessMapByStatement = HashMap::new();
-    ) -> (
+        let mut live_out: LivenessMapByStatement = HashMap::new();
-        HashMap<usize, HashSet<Rc<IrVariable>>>,
-        HashMap<usize, HashSet<Rc<IrVariable>>>,
-    ) {
-        let mut live_in: HashMap<usize, HashSet<Rc<IrVariable>>> = HashMap::new();
-        let mut live_out: HashMap<usize, HashSet<Rc<IrVariable>>> = HashMap::new();
 
         loop {
             let mut did_work = false;
@@ -73,17 +70,17 @@ impl IrBlock {
                 // this will need to be updated when we add jumps
                 if let Some(successor_live_in) = live_in.get(&(statement_index + 1)) {
                     let statement_live_out = live_out.get_mut(&statement_index).unwrap();
-                    for ir_variable in successor_live_in {
+                    for vr_variable in successor_live_in {
-                        if statement_live_out.insert(ir_variable.clone()) {
+                        if statement_live_out.insert(vr_variable.clone()) {
                             did_work = true;
                         }
                     }
                 }
 
                 // in: use(s) U ( out(s) - def(s) )
-                let mut new_ins = statement.uses();
+                let mut new_ins = statement.vr_uses();
                 let statement_live_out = live_out.get(&statement_index).unwrap();
-                let defs = statement.definitions();
+                let defs = statement.vr_definitions();
                 let rhs = statement_live_out - &defs;
                 new_ins.extend(rhs);
 
@@ -101,15 +98,20 @@ impl IrBlock {
         (live_in, live_out)
     }
 
-    fn interference_graph(&self) -> HashMap<Rc<IrVariable>, HashSet<Rc<IrVariable>>> {
+    fn interference_graph(
-        let mut all_variables: HashSet<Rc<IrVariable>> = HashSet::new();
+        &self,
+    ) -> HashMap<Rc<IrVirtualRegisterVariable>, HashSet<Rc<IrVirtualRegisterVariable>>> {
+        let mut all_vr_variables: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
         for statement in &self.statements {
-            all_variables.extend(statement.definitions());
+            all_vr_variables.extend(statement.vr_definitions());
-            all_variables.extend(statement.uses());
+            all_vr_variables.extend(statement.vr_uses());
         }
 
-        let mut graph: HashMap<Rc<IrVariable>, HashSet<Rc<IrVariable>>> = HashMap::new();
+        let mut graph: HashMap<
-        for variable in all_variables {
+            Rc<IrVirtualRegisterVariable>,
+            HashSet<Rc<IrVirtualRegisterVariable>>,
+        > = HashMap::new();
+        for variable in all_vr_variables {
             graph.insert(variable, HashSet::new());
         }
 
@@ -117,17 +119,18 @@ impl IrBlock {
 
         for (statement_index, statement) in self.statements.iter().enumerate() {
             let statement_live_out = live_out.get(&statement_index).unwrap();
-            for definition_variable in statement.definitions() {
+            for definition_vr_variable in statement.vr_definitions() {
                 for live_out_variable in statement_live_out {
-                    if definition_variable != *live_out_variable {
+                    // we do the following check to avoid adding an edge to itself
+                    if definition_vr_variable != *live_out_variable {
                         graph
-                            .get_mut(&definition_variable)
+                            .get_mut(&definition_vr_variable)
                             .unwrap()
                             .insert(live_out_variable.clone());
                         graph
                             .get_mut(live_out_variable)
                             .unwrap()
-                            .insert(definition_variable.clone());
+                            .insert(definition_vr_variable.clone());
                     }
                 }
             }
@@ -135,6 +138,23 @@ impl IrBlock {
 
         graph
     }
+
+    pub fn assign_registers(&mut self) {
+        let mut spills: HashSet<Rc<IrVirtualRegisterVariable>> = HashSet::new();
+        loop {
+            let (registers, new_spills) = register_assignment::assign_registers(self);
+            if spills != new_spills {
+                spills = new_spills;
+                // mutate all IrVirtualRegisters to constituent statements
+                for statement in &mut self.statements {
+                    statement.propagate_spills(&spills);
+                }
+            } else {
+                println!("{:?}", registers);
+                break;
+            }
+        }
+    }
 }
 
 impl Display for IrBlock {
@@ -154,8 +174,8 @@ impl Display for IrBlock {
             )?;
         }
         writeln!(f, "    // ---- {} meta ----", self.debug_label)?;
-        writeln!(f, "    // definitions: {:?}", self.definitions())?;
+        writeln!(f, "    // definitions: {:?}", self.vr_definitions())?;
-        writeln!(f, "    // uses: {:?}", self.uses())?;
+        writeln!(f, "    // uses: {:?}", self.vr_uses())?;
         writeln!(
             f,
             "    // interference graph: {:?}",
@@ -164,3 +184,168 @@ impl Display for IrBlock {
         Ok(())
     }
 }
+
+mod register_assignment {
+    use super::*;
+
+    #[derive(Debug)]
+    struct WorkItem {
+        vr: Rc<IrVirtualRegisterVariable>,
+        edges: HashSet<Rc<IrVirtualRegisterVariable>>,
+        color: bool,
+    }
+
+    pub fn assign_registers(
+        block: &IrBlock,
+    ) -> (
+        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 {
+                // 2. coloring by simplification
+                //    try to find a node (virtual register) with less than k outgoing edges,
+                //    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<
+                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();
+
+            while let Some(work_item) = work_stack.pop() {
+                if work_item.color {
+                    assign_register(&work_item, &mut rebuilt_graph, k, &mut register_assignments);
+                } else {
+                    // first, 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;
+                    for edge in &work_item.edges {
+                        if register_assignments.contains_key(edge) {
+                            number_of_assigned_edges += 1;
+                        }
+                    }
+
+                    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());
+                    }
+                }
+            }
+
+            if spills.eq(&new_spills) {
+                return (register_assignments, spills);
+            } else {
+                spills = new_spills;
+            }
+        }
+    }
+
+    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_call.rs

@@ -1,5 +1,6 @@
 use crate::ir::ir_expression::IrExpression;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::{IrStackVariable, IrVariable, IrVirtualRegisterVariable};
+use crate::type_info::TypeInfo;
 use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
@@ -17,13 +18,19 @@ impl IrCall {
         }
     }
 
-    pub fn uses(&self) -> HashSet<Rc<IrVariable>> {
+    pub fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
         let mut set = HashSet::new();
         for argument in &self.arguments {
-            set.extend(argument.uses());
+            set.extend(argument.vr_uses())
         }
         set
     }
+
+    pub fn propagate_spills(&mut self, spills: &HashSet<Rc<IrVirtualRegisterVariable>>) {
+        for argument in &mut self.arguments {
+            argument.propagate_spills(spills);
+        }
+    }
 }
 
 impl Display for IrCall {

dmc-lib/src/ir/ir_expression.rs

@@ -1,16 +1,59 @@
 use crate::ir::ir_parameter::IrParameter;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::{IrStackVariable, IrVariable, IrVirtualRegisterVariable};
 use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
 
 pub enum IrExpression {
     Parameter(Rc<IrParameter>),
-    Variable(Rc<IrVariable>),
+    Variable(IrVariable),
     Int(i32),
     String(Rc<str>),
 }
 
+impl IrExpression {
+    pub fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
+        match self {
+            IrExpression::Parameter(_) => HashSet::new(),
+            IrExpression::Variable(ir_variable) => {
+                let mut set = HashSet::new();
+                if let IrVariable::VirtualRegister(vr_variable) = ir_variable {
+                    set.insert(vr_variable.clone());
+                }
+                set
+            }
+            IrExpression::Int(_) => HashSet::new(),
+            IrExpression::String(_) => HashSet::new(),
+        }
+    }
+
+    pub fn propagate_spills(&mut self, spills: &HashSet<Rc<IrVirtualRegisterVariable>>) {
+        match self {
+            IrExpression::Parameter(_) => {
+                // no-op
+            }
+            IrExpression::Variable(ir_variable) => {
+                if let IrVariable::VirtualRegister(vr_variable) = ir_variable {
+                    if spills.contains(vr_variable) {
+                        let name = vr_variable.name().to_string();
+                        let type_info = vr_variable.type_info().clone();
+                        let _ = std::mem::replace(
+                            ir_variable,
+                            IrVariable::Stack(IrStackVariable::new(&name, type_info)),
+                        );
+                    }
+                }
+            }
+            IrExpression::Int(_) => {
+                // no-op
+            }
+            IrExpression::String(_) => {
+                // no-op
+            }
+        }
+    }
+}
+
 impl Display for IrExpression {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -29,24 +72,3 @@ impl Display for IrExpression {
         }
     }
 }
-
-impl IrExpression {
-    pub fn uses(&self) -> HashSet<Rc<IrVariable>> {
-        let mut set = HashSet::new();
-        match self {
-            IrExpression::Parameter(_) => {
-                // no-op
-            }
-            IrExpression::Variable(ir_variable) => {
-                set.insert(ir_variable.clone());
-            }
-            IrExpression::Int(_) => {
-                // no-op
-            }
-            IrExpression::String(_) => {
-                // no-op
-            }
-        }
-        set
-    }
-}

dmc-lib/src/ir/ir_function.rs

@@ -1,6 +1,5 @@
 use crate::ir::ir_block::IrBlock;
 use crate::ir::ir_parameter::IrParameter;
-use crate::ir::ir_variable::IrVariable;
 use crate::type_info::TypeInfo;
 use std::cell::RefCell;
 use std::fmt::Display;
@@ -27,6 +26,10 @@ impl IrFunction {
             entry,
         }
     }
+
+    pub fn assign_registers(&mut self) {
+        self.entry.borrow_mut().assign_registers();
+    }
 }
 
 impl Display for IrFunction {

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;
+use crate::ir::ir_variable::{IrVariable, IrVirtualRegisterVariable};
 use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
@@ -29,11 +29,11 @@ impl Display for IrOperation {
 }
 
 impl IrOperation {
-    pub fn uses(&self) -> HashSet<Rc<IrVariable>> {
+    pub fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
         match self {
-            IrOperation::Load(ir_expression) => ir_expression.uses(),
+            IrOperation::Load(ir_expression) => ir_expression.vr_uses(),
-            IrOperation::Add(ir_add) => ir_add.uses(),
+            IrOperation::Add(ir_add) => ir_add.vr_uses(),
-            IrOperation::Call(ir_call) => ir_call.uses(),
+            IrOperation::Call(ir_call) => ir_call.vr_uses(),
         }
     }
 }

dmc-lib/src/ir/ir_return.rs

@@ -1,5 +1,5 @@
 use crate::ir::ir_expression::IrExpression;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::IrVirtualRegisterVariable;
 use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
@@ -13,8 +13,14 @@ impl IrReturn {
         Self { value }
     }
 
-    pub fn uses(&self) -> HashSet<Rc<IrVariable>> {
+    pub fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
-        self.value.as_ref().map_or(HashSet::new(), |v| v.uses())
+        self.value.as_ref().map_or(HashSet::new(), |v| v.vr_uses())
+    }
+
+    pub fn propagate_spills(&mut self, spills: &HashSet<Rc<IrVirtualRegisterVariable>>) {
+        if let Some(ir_expression) = self.value.as_mut() {
+            ir_expression.propagate_spills(spills);
+        }
     }
 }
 

dmc-lib/src/ir/ir_statement.rs

@@ -1,7 +1,7 @@
 use crate::ir::ir_assign::IrAssign;
 use crate::ir::ir_call::IrCall;
 use crate::ir::ir_return::IrReturn;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::IrVirtualRegisterVariable;
 use std::collections::HashSet;
 use std::fmt::{Display, Formatter};
 use std::rc::Rc;
@@ -13,34 +13,34 @@ pub enum IrStatement {
 }
 
 impl IrStatement {
-    pub fn definitions(&self) -> HashSet<Rc<IrVariable>> {
+    pub fn vr_definitions(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
-        let mut set = HashSet::new();
         match self {
-            IrStatement::Assign(ir_assign) => {
+            IrStatement::Assign(ir_assign) => ir_assign.vr_definitions(),
-                set.insert(ir_assign.destination().clone());
+            IrStatement::Call(_) => HashSet::new(),
+            IrStatement::Return(_) => HashSet::new(),
         }
-            IrStatement::Call(_) => {
-                // no-op
-            }
-            IrStatement::Return(_) => {
-                // no-op
-            }
-        }
-        set
     }
 
-    pub fn uses(&self) -> HashSet<Rc<IrVariable>> {
+    pub fn vr_uses(&self) -> HashSet<Rc<IrVirtualRegisterVariable>> {
-        let mut set = HashSet::new();
+        match self {
+            IrStatement::Assign(ir_assign) => ir_assign.vr_uses(),
+            IrStatement::Call(ir_call) => ir_call.vr_uses(),
+            IrStatement::Return(ir_return) => ir_return.vr_uses(),
+        }
+    }
+
+    pub fn propagate_spills(&mut self, spills: &HashSet<Rc<IrVirtualRegisterVariable>>) {
         match self {
             IrStatement::Assign(ir_assign) => {
-                set.extend(ir_assign.uses());
+                ir_assign.propagate_spills(spills);
             }
             IrStatement::Call(ir_call) => {
-                set.extend(ir_call.uses());
+                ir_call.propagate_spills(spills);
+            }
+            IrStatement::Return(ir_return) => {
+                ir_return.propagate_spills(spills);
             }
-            IrStatement::Return(ir_return) => set.extend(ir_return.uses()),
         }
-        set
     }
 }
 

dmc-lib/src/ir/ir_variable.rs

@@ -3,13 +3,43 @@ use std::fmt::{Debug, Display, Formatter};
 use std::hash::{Hash, Hasher};
 use std::rc::Rc;
 
-/// Represents a virtual register
+pub enum IrVariable {
-pub struct IrVariable {
+    VirtualRegister(Rc<IrVirtualRegisterVariable>),
+    Stack(IrStackVariable),
+}
+
+impl IrVariable {
+    pub fn new(name: &str, type_info: TypeInfo) -> IrVariable {
+        IrVariable::VirtualRegister(Rc::new(IrVirtualRegisterVariable::new(name, type_info)))
+    }
+
+    pub fn type_info(&self) -> &TypeInfo {
+        match self {
+            IrVariable::VirtualRegister(vr_variable) => vr_variable.type_info(),
+            IrVariable::Stack(stack_variable) => stack_variable.type_info(),
+        }
+    }
+}
+
+impl Display for IrVariable {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            IrVariable::VirtualRegister(vr_variable) => {
+                write!(f, "{}", vr_variable)
+            }
+            IrVariable::Stack(stack_variable) => {
+                write!(f, "{}", stack_variable)
+            }
+        }
+    }
+}
+
+pub struct IrVirtualRegisterVariable {
     name: Rc<str>,
     type_info: TypeInfo,
 }
 
-impl IrVariable {
+impl IrVirtualRegisterVariable {
     pub fn new(name: &str, type_info: TypeInfo) -> Self {
         Self {
             name: name.into(),
@@ -26,28 +56,54 @@ impl IrVariable {
     }
 }
 
-impl Display for IrVariable {
+impl Display for IrVirtualRegisterVariable {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self.name)
     }
 }
 
-impl Debug for IrVariable {
+impl Debug for IrVirtualRegisterVariable {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         write!(f, "{}", self.name)
     }
 }
 
-impl Hash for IrVariable {
+impl Eq for IrVirtualRegisterVariable {}
-    fn hash<H: Hasher>(&self, state: &mut H) {
-        self.name.hash(state);
-    }
-}
 
-impl PartialEq for IrVariable {
+impl PartialEq for IrVirtualRegisterVariable {
     fn eq(&self, other: &Self) -> bool {
         self.name == other.name
     }
 }
 
-impl Eq for IrVariable {}
+impl Hash for IrVirtualRegisterVariable {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+    }
+}
+
+pub struct IrStackVariable {
+    name: Rc<str>,
+    type_info: TypeInfo,
+    offset: Option<isize>,
+}
+
+impl IrStackVariable {
+    pub fn new(name: &str, type_info: TypeInfo) -> Self {
+        Self {
+            name: name.into(),
+            type_info,
+            offset: None,
+        }
+    }
+
+    pub fn type_info(&self) -> &TypeInfo {
+        &self.type_info
+    }
+}
+
+impl Display for IrStackVariable {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.name)
+    }
+}

dmc-lib/src/symbol.rs

@@ -1,6 +1,6 @@
 use crate::ir::ir_expression::IrExpression;
 use crate::ir::ir_parameter::IrParameter;
-use crate::ir::ir_variable::IrVariable;
+use crate::ir::ir_variable::{IrVariable, IrVirtualRegisterVariable};
 use crate::type_info::TypeInfo;
 use std::cell::RefCell;
 use std::rc::Rc;
@@ -100,7 +100,7 @@ impl ParameterSymbol {
 pub struct VariableSymbol {
     name: Rc<str>,
     type_info: Option<TypeInfo>,
-    ir_variable: Option<Rc<IrVariable>>,
+    vr_variable: Option<Rc<IrVirtualRegisterVariable>>,
 
     #[deprecated]
     register: Option<usize>,
@@ -111,7 +111,7 @@ impl VariableSymbol {
         Self {
             name: name.into(),
             type_info: None,
-            ir_variable: None,
+            vr_variable: None,
             register: None,
         }
     }
@@ -134,12 +134,12 @@ impl VariableSymbol {
             .expect("TypeInfo not initialized. Did you type check?")
     }
 
-    pub fn set_ir_variable(&mut self, ir_variable: Rc<IrVariable>) {
+    pub fn set_vr_variable(&mut self, ir_variable: Rc<IrVirtualRegisterVariable>) {
-        self.ir_variable = Some(ir_variable);
+        self.vr_variable = Some(ir_variable);
     }
 
-    pub fn ir_variable(&self) -> &Rc<IrVariable> {
+    pub fn vr_variable(&self) -> &Rc<IrVirtualRegisterVariable> {
-        self.ir_variable
+        self.vr_variable
             .as_ref()
             .expect("ir_variable not yet initialized")
     }
@@ -168,9 +168,9 @@ impl ExpressibleSymbol {
             ExpressibleSymbol::Parameter(parameter_symbol) => {
                 IrExpression::Parameter(parameter_symbol.borrow().ir_parameter().clone())
             }
-            ExpressibleSymbol::Variable(variable_symbol) => {
+            ExpressibleSymbol::Variable(variable_symbol) => IrExpression::Variable(
-                IrExpression::Variable(variable_symbol.borrow().ir_variable().clone())
+                IrVariable::VirtualRegister(variable_symbol.borrow().vr_variable().clone()),
-            }
+            ),
         }
     }
 }

examples/register_alloc_test.dm

@@ -0,0 +1,8 @@
+extern fn println(message: Any) -> Void
+
+fn main()
+    let a = 1
+    let b = 2
+    let c = 3
+    let x = a + b + c
+end