Reid/src/compiler.rs

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

use super::errors::CompilerError;
use super::parser::{Expression, LiteralPattern, ParsedReid, Pattern, Position, Statement};
use super::vm::{FunctionSignature, VariableType};

type Variable = (HeapID, VariableType);
pub type FuncID = usize;
pub type HeapID = usize;
pub type RegID = usize;

#[derive(Debug, Clone)]
pub enum Command {
    InitializeVariable(HeapID, VariableType), // Initializes new variable to HeapID at VariableType
    BeginScope,                               // Begins new Scope
    EndScope,                                 // Ends Scope
    Pop(RegID),                               // Pop into registery at RegID
    Push(RegID),                              // Push out of registery at RegID
    AssignVariable(HeapID, RegID),            // Assign variable from registery at RegID
    VarToReg(HeapID, RegID),                  // Bring Variable to registery at RegID
    StringLit(String),                        // Bring String Literal to Stack
    I32Lit(i32),                              // Bring i32 Literal to Stack
    FunctionCall(FuncID),                     // Call Function at FuncID
}

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

#[derive(Debug)]
pub struct CompiledReid {
    pub list: Vec<Command>,
}

impl Compiler {
    pub fn from(parsed: ParsedReid) -> Compiler {
        Compiler {
            parsed,
            root_scope: Scope::default(),
            list: Vec::new(),
        }
    }

    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::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_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<(usize, &FunctionSignature)> {
        let mut found = None;
        for (idx, func) in self.functions.iter().enumerate() {
            if func == signature {
                found = Some((idx, 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> {
        if self.variables.contains_key(&variable) {
            Err(CompilerError::VariableExists(pos, variable))
        } else if let Some(inner) = &mut self.inner_scope {
            inner.add_var(pos, variable, vtype)
        } else {
            self.variables.insert(variable, (self.counter, vtype));
            self.counter += 1;
            Ok(())
        }
    }

    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))
        }
    }
}