Reid/src/compiler.rs

use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::num::ParseIntError;

use super::errors::CompilerError;
use super::parser::{
    ArithmeticExpression, Expression, LiteralPattern, ParsedReid, Pattern, Statement,
};
use super::vm::VariableType::*;
use super::vm::{Command, CompiledReid, FuncID, FunctionSignature, HeapID, Position, VariableType};

pub type Variable = (HeapID, VariableType);

#[derive(Debug, Clone)]
pub enum ArithmeticOp {
    Add(VariableType, VariableType, VariableType),
    Subtract(VariableType, VariableType, VariableType),
    Mult(VariableType, VariableType, VariableType),
    Pow(VariableType, VariableType, VariableType),
    Div(VariableType, VariableType, VariableType),
    Mod(VariableType, VariableType, VariableType),
}

pub struct Compiler {
    parsed: ParsedReid,
    root_scope: Scope,
    list: Vec<Command>,
    possible_arit_possib: Vec<ArithmeticOp>,
}

impl Compiler {
    pub fn from(parsed: ParsedReid) -> Compiler {
        let list = vec![
            ArithmeticOp::Add(TypeI32, TypeI32, TypeI32),
            ArithmeticOp::Subtract(TypeI32, TypeI32, TypeI32),
            ArithmeticOp::Mult(TypeI32, TypeI32, TypeI32),
            ArithmeticOp::Pow(TypeI32, TypeI32, TypeI32),
            ArithmeticOp::Div(TypeI32, TypeI32, TypeI32),
            ArithmeticOp::Mod(TypeI32, TypeI32, TypeI32),
            ArithmeticOp::Add(TypeString, TypeI32, TypeString),
            ArithmeticOp::Add(TypeString, TypeString, TypeString),
        ];
        Compiler {
            parsed,
            root_scope: Scope::default(),
            list: Vec::new(),
            possible_arit_possib: list,
        }
    }

    pub fn with_builtin_functions<T: Into<Vec<FunctionSignature>>>(mut self, list: T) -> Compiler {
        self.add_builtin_functions(list.into());
        self
    }

    pub fn compile(mut self) -> Result<CompiledReid, CompilerError> {
        self.handle_expression(self.parsed.0.clone())?;
        Ok(CompiledReid { list: self.list })
    }

    pub fn add_builtin_functions<T: Into<Vec<FunctionSignature>>>(&mut self, list: T) {
        for func in list.into() {
            self.root_scope.add_builtin_function(func);
        }
    }

    fn handle_expression(
        &mut self,
        exp: Expression,
    ) -> Result<Option<VariableType>, CompilerError> {
        match exp {
            Expression::BlockExpr(pos, list) => {
                self.list.push(Command::BeginScope);
                self.root_scope.begin_scope();

                for statement in list {
                    self.handle_statement(statement)?;
                }

                self.root_scope.end_scope(pos)?;
                self.list.push(Command::EndScope);
                Ok(None)
            }
            Expression::FunctionCall(pos, name, args) => {
                let mut arguments = Vec::new();

                for expr in args {
                    if let Some(vtype) = self.handle_expression(expr)? {
                        arguments.push(vtype);
                    } else {
                        return Err(CompilerError::CanNotAssignVoidType);
                    }
                }
                if let Some(func) = self
                    .root_scope
                    .find_function(&FunctionSignature::new(name.clone(), arguments.clone()))
                {
                    self.list.push(Command::FunctionCall(func.0));
                    Ok(None)
                } else {
                    Err(CompilerError::FunctionNotFound(pos, name, arguments))
                }
            }
            Expression::ArithmeticExpression(pos, val) => {
                self.handle_arithmetic_expression(pos, *val)
            }
            Expression::ValueRef(_, val) => match val {
                Pattern::IdentPattern(pos, ident) => {
                    if let Some(var) = self.root_scope.get(ident.clone()) {
                        self.list.push(Command::VarToReg(var.0, 0));
                        self.list.push(Command::Push(0));
                        Ok(Some(var.1))
                    } else {
                        Err(CompilerError::VariableNotExists(pos, ident))
                    }
                }
                Pattern::LiteralPattern(_, literal) => {
                    let vtype = self.handle_literal(literal)?;
                    Ok(Some(vtype))
                }
            },
        }
    }

    fn handle_arithmetic_expression(
        &mut self,
        pos: Position,
        exp: ArithmeticExpression,
    ) -> Result<Option<VariableType>, CompilerError> {
        let (exp1, exp2) = match exp.clone() {
            ArithmeticExpression::Add(e1, e2) => (e1, e2),
            ArithmeticExpression::Subtract(e1, e2) => (e1, e2),
            ArithmeticExpression::Mult(e1, e2) => (e1, e2),
            ArithmeticExpression::Pow(e1, e2) => (e1, e2),
            ArithmeticExpression::Div(e1, e2) => (e1, e2),
            ArithmeticExpression::Mod(e1, e2) => (e1, e2),
        };
        let type1 = self.handle_expression(exp1)?;
        let type2 = self.handle_expression(exp2)?;
        if let (Some(type1), Some(type2)) = (type1, type2) {
            let mut return_type = None;
            for op in &self.possible_arit_possib {
                match op {
                    ArithmeticOp::Add(t1, t2, t3) => {
                        if let ArithmeticExpression::Add(..) = exp {
                            if (t1 == &type1 && t2 == &type2) || (t2 == &type1 && t1 == &type2) {
                                return_type = Some(t3);
                                break;
                            }
                        }
                    }
                    ArithmeticOp::Subtract(t1, t2, t3) => {
                        if let ArithmeticExpression::Subtract(..) = exp {
                            if (t1 == &type1 && t2 == &type2) || (t2 == &type1 && t1 == &type2) {
                                return_type = Some(t3);
                                break;
                            }
                        }
                    }
                    ArithmeticOp::Mult(t1, t2, t3) => {
                        if let ArithmeticExpression::Mult(..) = exp {
                            if (t1 == &type1 && t2 == &type2) || (t2 == &type1 && t1 == &type2) {
                                return_type = Some(t3);
                                break;
                            }
                        }
                    }
                    ArithmeticOp::Div(t1, t2, t3) => {
                        if let ArithmeticExpression::Div(..) = exp {
                            if (t1 == &type1 && t2 == &type2) || (t2 == &type1 && t1 == &type2) {
                                return_type = Some(t3);
                                break;
                            }
                        }
                    }
                    ArithmeticOp::Pow(t1, t2, t3) => {
                        if let ArithmeticExpression::Pow(..) = exp {
                            if (t1 == &type1 && t2 == &type2) || (t2 == &type1 && t1 == &type2) {
                                return_type = Some(t3);
                                break;
                            }
                        }
                    }
                    ArithmeticOp::Mod(t1, t2, t3) => {
                        if let ArithmeticExpression::Mod(..) = exp {
                            if (t1 == &type1 && t2 == &type2) || (t2 == &type1 && t1 == &type2) {
                                return_type = Some(t3);
                                break;
                            }
                        }
                    }
                }
            }
            let (command, error) = match exp {
                ArithmeticExpression::Add(..) => {
                    (Command::Add, ArithmeticOp::Add(type1, type2, TypeI32))
                }
                ArithmeticExpression::Subtract(..) => (
                    Command::Subtract,
                    ArithmeticOp::Subtract(type1, type2, TypeI32),
                ),
                ArithmeticExpression::Mult(..) => {
                    (Command::Mult, ArithmeticOp::Mult(type1, type2, TypeI32))
                }
                ArithmeticExpression::Pow(..) => {
                    (Command::Pow, ArithmeticOp::Pow(type1, type2, TypeI32))
                }
                ArithmeticExpression::Div(..) => {
                    (Command::Div, ArithmeticOp::Div(type1, type2, TypeI32))
                }
                ArithmeticExpression::Mod(..) => {
                    (Command::Mod, ArithmeticOp::Mod(type1, type2, TypeI32))
                }
            };
            if let Some(t) = return_type {
                self.list.push(command);
                Ok(Some(*t))
            } else {
                Err(CompilerError::ArithmeticExpressionFailed(
                    pos,
                    Box::new(CompilerError::InvalidArithmeticOperation(error)),
                ))
            }
        } else {
            Err(CompilerError::ArithmeticExpressionFailed(
                pos,
                Box::new(CompilerError::CanNotAssignVoidType),
            ))
        }
    }

    fn handle_statement(
        &mut self,
        statement: Statement,
    ) -> Result<Option<VariableType>, CompilerError> {
        match statement {
            Statement::LetStatement(pos, ident, val) => {
                let res = self.handle_expression(*val);
                match res {
                    Ok(vtype) => {
                        if let Some(vtype) = vtype {
                            self.root_scope.add_var(pos, ident.clone(), vtype)?;
                            let var = self.root_scope.get(ident).unwrap();
                            self.list.push(Command::InitializeVariable(var.0, var.1));
                            self.list.push(Command::Pop(0));
                            self.list.push(Command::AssignVariable(var.0, 0));
                            Ok(None)
                        } else {
                            Err(CompilerError::LetFailed(
                                pos,
                                Box::new(CompilerError::CanNotAssignVoidType),
                            ))
                        }
                    }
                    Err(err) => Err(CompilerError::LetFailed(pos, Box::new(err))),
                }
            }
            Statement::ExprStatement(_, expr) => self.handle_expression(expr),
        }
    }

    fn handle_literal(&mut self, pattern: LiteralPattern) -> Result<VariableType, CompilerError> {
        match pattern {
            LiteralPattern::StringLit(string) => {
                self.list.push(Command::StringLit(string));
                Ok(VariableType::TypeString)
            }
            LiteralPattern::Integer32Lit(string) => {
                self.list.push(Command::I32Lit(Compiler::handle_parseint(
                    string.parse::<i32>(),
                )?));
                Ok(VariableType::TypeI32)
            }
        }
    }

    fn handle_parseint<T>(res: Result<T, ParseIntError>) -> Result<T, CompilerError> {
        match res {
            Ok(i) => Ok(i),
            Err(err) => Err(CompilerError::ParseIntError(err)),
        }
    }
}

#[derive(Default)]
pub struct Scope {
    counter: HeapID,
    variables: HashMap<String, Variable>,
    functions: Vec<FunctionSignature>,
    inner_scope: Option<Box<Scope>>,
}

impl Scope {
    fn get(&self, variable: String) -> Option<Variable> {
        if let Some(val) = self.variables.get(&variable) {
            Some(*val)
        } else if let Some(inner) = &self.inner_scope {
            if let Some(val) = inner.get(variable) {
                Some((val.0 + self.counter, val.1))
            } else {
                None
            }
        } else {
            None
        }
    }

    fn find_function(&self, signature: &FunctionSignature) -> Option<(FuncID, &FunctionSignature)> {
        let mut found = None;
        for (idx, func) in self.functions.iter().enumerate() {
            if func == signature {
                found = Some((idx as u16, func));
            }
        }
        found
    }

    fn add_builtin_function(&mut self, function: FunctionSignature) -> bool {
        if self.find_function(&function).is_some() {
            false
        } else {
            self.functions.push(function);
            true
        }
    }

    fn add_var(
        &mut self,
        pos: Position,
        variable: String,
        vtype: VariableType,
    ) -> Result<(), CompilerError> {
        let entry = self.variables.entry(variable.clone());
        if let Entry::Vacant(e) = entry {
            if let Some(inner) = &mut self.inner_scope {
                inner.add_var(pos, variable, vtype)
            } else {
                e.insert((self.counter, vtype));
                self.counter += 1;
                Ok(())
            }
        } else {
            Err(CompilerError::VariableExists(pos, variable))
        }
    }

    fn begin_scope(&mut self) {
        if let Some(inner) = &mut self.inner_scope {
            inner.begin_scope();
        } else {
            self.inner_scope = Some(Box::default());
        }
    }

    fn end_scope(&mut self, pos: Position) -> Result<(), CompilerError> {
        if let Some(inner) = &mut self.inner_scope {
            if inner.inner_scope.is_some() {
                inner.end_scope(pos)?;
            } else {
                self.inner_scope = None;
            }
            Ok(())
        } else {
            Err(CompilerError::InvalidScopeExit(pos))
        }
    }
}