mod.rs

use std::collections::HashMap;

use crate::{ast::model::*, lex::{UnaryOperator, BinaryOperator}, utility::{generate_clause, generate_end, SyntaxError}};

static FUNCTION_PROLOGUE_START: &str = "push\t%rbp\n\tmov\t%rsp, %rbp";
static FUNCTION_PROLOGUE_END: &str = "mov\t%rbp, %rsp\n\tpop\t%rbp\n\tret\n";

pub struct CodeGenerator {
  symbol_table: HashMap<String, i64>,
  stack_index: i64,
}

impl CodeGenerator {
  pub fn new() -> Self {
    Self { symbol_table: HashMap::new(), stack_index: 0 }
  }
  pub fn generate(&mut self, root: &Program) -> Result<String, SyntaxError> {
    self.generate_function(&root.func)
  }
  fn generate_function(&mut self, func: &Function) -> Result<String, SyntaxError> {
    let mut body = vec![];
    let mut return_flag = false;
    for x in &func.body {
      let mut st = self.generate_statement(x)?;
      st = st.replace("\n", "\n\t");

      match x {
        Statement::Return(_) => {
          st = format!("{}\n\t{}", st, FUNCTION_PROLOGUE_END);
          return_flag = true;
        },
        _ => {}
      }
      body.push(st);
    }

    let end = if return_flag {format!("")} else {format!("mov\t$0, %rax\n\t{}", FUNCTION_PROLOGUE_END)};

    Ok(format!("
\t.globl {name}
{name}:
\t{start}
\t{body_text}
\t{end}
", start = FUNCTION_PROLOGUE_START, name = func.name, body_text = body.join("\n\t")))
  }
  fn generate_statement(&mut self, st: &Statement) -> Result<String, SyntaxError> {
    match st {
      Statement::Return(val) => {
        let exp = self.generate_expression(val)?;

        Ok(format!("{}", exp))
      },
      Statement::Declare { name, exp } => {
        if self.symbol_table.contains_key(name.as_ref()) {
          Err(SyntaxError::new_codegen_error(format!("re-declaration of variable {}", name).to_string()))
        }
        else {
          let mut assembly_exp = self.generate_expression(&Expression::Const(0))?;
          if let Some(exp_some) = exp {
            assembly_exp = self.generate_expression(exp_some)?;
          };
          self.stack_index -= 8;
          self.symbol_table.insert(name.to_string(), self.stack_index);
          Ok(format!("{}\npush\t%rax", assembly_exp))
        }
      },
      Statement::Exp(val) => self.generate_expression(val)
    }
  }
  fn generate_expression(&self, exp: &Expression) -> Result<String, SyntaxError> {
    match exp {
      Expression::Const(val) => Ok(format!("mov\t${}, %rax", val.to_string())),
      Expression::Unary { op, exp } => {
        let inner_exp = self.generate_expression(exp)?;
        let ext_exp = match op {
          UnaryOperator::Negation => format!("neg\t%rax"),
          UnaryOperator::BitwiseComplement => format!("not\t%rax"),
          UnaryOperator::LogicalNegation => {
            format!("cmp\t$0, %rax\nmov\t$0, %rax\nsete\t%al")
          }
        };
        Ok(format!("{}\n{}", inner_exp, ext_exp))
      },
      Expression::Binary { exp1, op, exp2 } => {
        let exp1 = self.generate_expression(exp1)?;
        let exp2 = self.generate_expression(exp2)?;
        let inner_exp = format!("{}\npush\t%rax\n{}\npop\t%rcx", exp1, exp2);
        let ext_exp = match op {
          BinaryOperator::Addition => format!("add\t%rcx, %rax"),
          BinaryOperator::Multiplication => format!("imul\t%rcx, %rax"),
          BinaryOperator::Minus => format!("sub\t%rax, %rcx\nmov\t%rcx, %rax"),
          BinaryOperator::Division => format!("mov\t%rax, %rbx\nmov\t%rcx, %rax\ncqo\nidiv\t%rbx"),
          BinaryOperator::Equal => format!("cmp\t%rax, %rcx\nsete\t%al"),
          BinaryOperator::NotEqual => format!("cmp\t%rax, %rcx\ncmp\t$0, %rax\nsetne\t%al"),
          BinaryOperator::LessThan => format!("cmp\t%rax, %rcx\nsetl\t%al"),
          BinaryOperator::LessThanOrEqual => format!("cmp\t%rax, %rcx\nsetle\t%al"),
          BinaryOperator::GreaterThan => format!("cmp\t%rax, %rcx\nsetg\t%al"),
          BinaryOperator::GreaterThanOrEqual => format!("cmp\t%rax, %rcx\nsetge\t%al"),
          BinaryOperator::And => {
            return Ok(format!("{}\ncmp\t$0, %rax\njne\t{_clause2}\njmp\t{_end}\n{_clause2}:\n{}\ncmp\t$0, %rax\ncmp $0, %rax\nsetne\t%al\n{_end}:", exp1, exp2, _clause2 = generate_clause(), _end = generate_end()));
          },
          BinaryOperator::Or => {
            return Ok(format!("{}\ncmp\t$0, %rax\nje\t{_clause2}\nmov\t$1, %rax\njmp\t{_end}\n{_clause2}:\n{}\ncmp\t$0, %rax\ncmp $0, %rax\nsetne\t%al\n{_end}:", exp1, exp2, _clause2 = generate_clause(), _end = generate_end()));
          },
        };
        Ok(format!("{}\n{}", inner_exp, ext_exp))
      },
      Expression::Assign { name, exp } => {
        match self.symbol_table.get(name.as_ref()) {
          None => Err(SyntaxError::new_codegen_error(format!("variable not declared {}", name).to_string())),
          Some(&offset) => {
            let assign_exp = self.generate_expression(exp)?;
            Ok(format!("{}\nmov\t%rax, {}(%rbp)", assign_exp, offset))
          }
        }
      },
      Expression::Var { name } => {
        match self.symbol_table.get(name.as_ref()) {
          None => Err(SyntaxError::new_codegen_error(format!("variable not declared {}", name).to_string())),
          Some(&offset) => {
            Ok(format!("mov\t{}(%rbp), %rax", offset))
          }
        }
      }
    }
  }
}