Add unary ops and if-then-else to IR Generator

src/compiler.rs

@@ -29,6 +29,7 @@ pub fn compile(code: &str) -> Vec<IrInstruction> {
 pub fn start_compiler() {
     let lines = io::stdin().lines();
     for line in lines.map_while(Result::ok) {
+        println!();
         for instruction in compile(&line) {
             println!("{instruction}");
         }

src/compiler/ir.rs

@@ -52,11 +52,11 @@ impl IrVar {
 
 impl fmt::Display for IrInstruction {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "{}           # From {}", self.instruction, self.loc)
+        write!(f, "{}", self.instruction)
     }
 }
 
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Clone, Eq, Hash)]
 pub struct IrInstruction {
     pub loc: CodeLocation,
     pub instruction: IrInstructionType,
@@ -68,7 +68,7 @@ impl IrInstruction {
     }
 }
 
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Clone, Eq, Hash)]
 pub enum IrInstructionType {
     LoadBoolConst(bool, IrVar),
     LoadIntConst(i64, IrVar),

src/compiler/ir_generator.rs

@@ -1,9 +1,14 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use crate::compiler::{
     ast::{AstNode, Expression::*},
-    ir::{IrInstruction, IrInstructionType::*, IrVar},
+    ir::{
+        IrInstruction,
+        IrInstructionType::{self, *},
+        IrVar,
+    },
     symtab::SymTab,
+    token::CodeLocation,
     variable::Type,
 };
 
@@ -11,31 +16,40 @@ pub fn generate_ir(ast: &AstNode) -> Vec<IrInstruction> {
     let mut instructions = Vec::new();
 
     let mut symbols = SymTab::new();
+    let mut labels = HashSet::new();
     let global_types = IrVar::new_global_types();
     let mut types = global_types.clone();
     for var in global_types.keys() {
         symbols.insert(&var.name, var.clone());
     }
 
-    let result = visit_ast_node(ast, &mut types, &mut symbols, &mut instructions);
+    let result = visit_ast_node(
+        ast,
+        &mut types,
+        &mut symbols,
+        &mut instructions,
+        &mut labels,
+    );
 
     match types.get(&result) {
         Some(Type::Int) => {
             let loc = instructions.last().unwrap().loc;
             let fn_var = symbols.get("print_int").clone();
+            let result_var = add_var(&Type::Bool, &mut types);
 
             instructions.push(IrInstruction::new(
                 loc,
-                Call(fn_var, vec![result], symbols.get("unit").clone()),
+                Call(fn_var, vec![result], result_var),
             ));
         }
         Some(Type::Bool) => {
             let loc = instructions.last().unwrap().loc;
             let fn_var = symbols.get("print_bool").clone();
+            let result_var = add_var(&Type::Bool, &mut types);
 
             instructions.push(IrInstruction::new(
                 loc,
-                Call(fn_var, vec![result], symbols.get("unit").clone()),
+                Call(fn_var, vec![result], result_var),
             ));
         }
         _ => (),
@@ -46,30 +60,41 @@ pub fn generate_ir(ast: &AstNode) -> Vec<IrInstruction> {
 
 fn add_var(var_type: &Type, types: &mut HashMap<IrVar, Type>) -> IrVar {
     let mut i = 1;
-    //let type_str = match var_type {
+    let mut var = IrVar::new(&format!("x{}", i));
-    //    Type::Int => "i",
-    //    Type::Bool => "b",
-    //    Type::Func(_, _) => "f",
-    //    Type::Unit => "u",
-    //};
-    let type_str = "x";
-
-    let mut var = IrVar::new(&format!("{}{}", type_str, i));
 
     while types.contains_key(&var) {
         i += 1;
-        var = IrVar::new(&format!("{}{}", type_str, i));
+        var = IrVar::new(&format!("x{}", i));
     }
 
     types.insert(var.clone(), var_type.clone());
     var
 }
 
+fn add_label(
+    label: &str,
+    loc: CodeLocation,
+    labels: &mut HashSet<IrInstructionType>,
+) -> IrInstruction {
+    let mut i = 1;
+
+    let mut instruction = IrInstructionType::Label(format!("{}{}", label, i));
+
+    while labels.contains(&instruction) {
+        i += 1;
+        instruction = IrInstructionType::Label(format!("{}{}", label, i));
+    }
+
+    labels.insert(instruction.clone());
+    IrInstruction::new(loc, instruction)
+}
+
 fn visit_ast_node(
     ast: &AstNode,
     types: &mut HashMap<IrVar, Type>,
     symbols: &mut SymTab<IrVar>,
     instructions: &mut Vec<IrInstruction>,
+    labels: &mut HashSet<IrInstructionType>,
 ) -> IrVar {
     match &ast.expr {
         EmptyLiteral() => symbols.get("unit").clone(),
@@ -87,15 +112,26 @@ fn visit_ast_node(
             var
         }
         Identifier(name) => symbols.get(name).clone(),
-        UnaryOp(_, _) => todo!(),
+        UnaryOp(op, expr) => {
+            let op_var = symbols.get(&format!("unary_{op}")).clone();
+            let expr_var = visit_ast_node(expr, types, symbols, instructions, labels);
+            let result_var = add_var(&ast.node_type, types);
+
+            instructions.push(IrInstruction::new(
+                ast.loc,
+                Call(op_var, vec![expr_var], result_var.clone()),
+            ));
+
+            result_var
+        }
         BinaryOp(left, op, right) => match *op {
-            "=" => todo!(),   // TODO Special handling
+            "=" => todo!(),
-            "and" => todo!(), // TODO Special handling
+            "and" => todo!(),
-            "or" => todo!(),  // TODO Special handling
+            "or" => todo!(),
             _ => {
                 let op_var = symbols.get(op).clone();
-                let left_var = visit_ast_node(left, types, symbols, instructions);
+                let left_var = visit_ast_node(left, types, symbols, instructions, labels);
-                let right_var = visit_ast_node(right, types, symbols, instructions);
+                let right_var = visit_ast_node(right, types, symbols, instructions, labels);
                 let result_var = add_var(&ast.node_type, types);
 
                 instructions.push(IrInstruction::new(
@@ -107,7 +143,59 @@ fn visit_ast_node(
             }
         },
         VarDeclaration(_, _, _) => todo!(),
-        Conditional(_, _, _) => todo!(),
+        Conditional(condition_expr, then_expr, else_expr) => match else_expr {
+            Some(else_expr) => {
+                let l_then = add_label("then", then_expr.loc, labels);
+                let l_else = add_label("else", else_expr.loc, labels);
+                let l_end = add_label("if_end", else_expr.loc, labels);
+
+                let cond_var = visit_ast_node(condition_expr, types, symbols, instructions, labels);
+                let result_var = add_var(&ast.node_type, types);
+
+                instructions.push(IrInstruction::new(
+                    condition_expr.loc,
+                    CondJump(cond_var, Box::new(l_then.clone()), Box::new(l_else.clone())),
+                ));
+
+                instructions.push(l_then);
+                let then_var = visit_ast_node(then_expr, types, symbols, instructions, labels);
+                instructions.push(IrInstruction::new(
+                    else_expr.loc,
+                    Copy(then_var, result_var.clone()),
+                ));
+                instructions.push(IrInstruction::new(
+                    else_expr.loc,
+                    Jump(Box::new(l_end.clone())),
+                ));
+
+                instructions.push(l_else);
+                let else_var = visit_ast_node(else_expr, types, symbols, instructions, labels);
+                instructions.push(IrInstruction::new(
+                    else_expr.loc,
+                    Copy(else_var, result_var.clone()),
+                ));
+                instructions.push(l_end);
+
+                result_var
+            }
+            None => {
+                let l_then = add_label("then", then_expr.loc, labels);
+                let l_end = add_label("if_end", then_expr.loc, labels);
+
+                let cond_var = visit_ast_node(condition_expr, types, symbols, instructions, labels);
+
+                instructions.push(IrInstruction::new(
+                    condition_expr.loc,
+                    CondJump(cond_var, Box::new(l_then.clone()), Box::new(l_end.clone())),
+                ));
+
+                instructions.push(l_then);
+                visit_ast_node(then_expr, types, symbols, instructions, labels);
+                instructions.push(l_end);
+
+                symbols.get("unit").clone()
+            }
+        },
         While(_, _) => todo!(),
         FunCall(_, _) => todo!(),
         Block(_) => todo!(),

src/compiler/token.rs

@@ -1,6 +1,7 @@
 use std::fmt;
 
-#[derive(Debug, Copy, Clone)]
+#[allow(clippy::derived_hash_with_manual_eq)]
+#[derive(Debug, Copy, Clone, Eq, Hash)]
 pub struct CodeLocation {
     line: usize,
     char: usize,