luallaby 0.1.0

use std::rc::Rc;

use crate::ast::{Exp, Func, PrefixExp};
use crate::compiler::{Compiler, LocalType, ScopeType};
use crate::error::Result;
use crate::lexer::Pos;
use crate::vm::{FuncObject, Literal, OpCode};
use crate::LuaError;

impl Compiler {
    pub(super) fn compile_func(&mut self, func: Func) -> Result<usize> {
        self.scope_enter(ScopeType::Func);
        self.code_push();
        self.scopes.declare_params(func.params);

        // Compile into new chunk
        self.compile_block(*func.block)?;

        // Ensure that function has at least one return
        let ret_reg = self.scopes.reg_reserve();
        self.code.emit(
            OpCode::Lit {
                val: Literal::Empty,
                dst_reg: ret_reg,
            },
            func.pos_end,
        );
        self.code.emit(OpCode::Return { ret_reg }, func.pos_end);
        self.scopes.reg_free(ret_reg);

        let scope = self
            .scope_leave(ScopeType::Func, self.code.get_pos_last())?
            .unwrap();
        let code = self.code_pop();

        // XXX: Does not need to be enforced, but Lua suite expects it
        if scope.ups.len() > 255 {
            return err!(LuaError::SyntaxLimit(
                "upvalues",
                Some(func.pos_start.line())
            ));
        }

        let def = FuncObject {
            regs: scope.regs.count(),
            locals_cap: scope.locals_cap,
            params: scope
                .locals
                .into_iter()
                .map(|loc| {
                    assert!(matches!(loc.typ, LocalType::Param));
                    loc.name
                })
                .collect(),
            varargs: func.varargs,
            ups: scope.ups.into_iter().map(|up| up.upper).collect(),
            code: Rc::new(code),
            linedefined: func.pos_start.line(),
            lastlinedefined: func.pos_end.line(),
        };

        let dst_reg = self.scopes.reg_reserve();
        self.code.emit(
            OpCode::Closure {
                func: Box::new(def),
                dst_reg,
            },
            func.pos_end,
        );
        Ok(dst_reg)
    }

    pub(super) fn compile_call(
        &mut self,
        pos_call: Pos,
        prefix: PrefixExp,
        mut args: Vec<Exp>,
        tail: bool,
    ) -> Result<usize> {
        // Reserve return register first, so it is more likely that
        // its value gets dropped at next use of register, increasing
        // likelyhood that object gets collected
        let ret_reg = self.scopes.reg_reserve();

        let args_reg = self.scopes.reg_reserve();
        self.code.emit(
            OpCode::Lit {
                val: Literal::EmptyCap(args.len()),
                dst_reg: args_reg,
            },
            pos_call,
        );
        let last_arg = args.pop();
        for arg in args.into_iter() {
            let arg_reg = self.compile_exp(arg)?;
            self.code.emit(
                OpCode::Append {
                    src_reg: arg_reg,
                    dst_reg: args_reg,
                },
                pos_call,
            );
            self.scopes.reg_free(arg_reg);
        }
        if let Some(arg) = last_arg {
            let arg_reg = self.compile_exp(arg)?;
            self.code.emit(
                OpCode::Extend {
                    src_reg: arg_reg,
                    dst_reg: args_reg,
                },
                pos_call,
            );
            self.scopes.reg_free(arg_reg);
        }

        let func_reg = self.compile_prefix_exp(prefix)?;

        self.code.emit(
            OpCode::Call {
                func_reg,
                args_reg,
                ret_reg,
                tail,
            },
            pos_call,
        );

        self.scopes.reg_free(args_reg);
        self.scopes.reg_free(func_reg);

        Ok(ret_reg)
    }

    pub(super) fn compile_call_method(
        &mut self,
        pos_call: Pos,
        prefix: PrefixExp,
        name: String,
        mut args: Vec<Exp>,
        tail: bool,
    ) -> Result<usize> {
        // Reserve return register first, so it is more likely that
        // its value gets dropped at next use of register, increasing
        // likelyhood that object gets collected
        let ret_reg = self.scopes.reg_reserve();

        let obj_reg = self.compile_prefix_exp(prefix)?;
        let func_reg = self.scopes.reg_reserve();
        self.code.emit(
            OpCode::Copy {
                src_reg: obj_reg,
                dst_reg: func_reg,
            },
            pos_call,
        );

        let args_reg = self.scopes.reg_reserve();
        self.code.emit(
            OpCode::Lit {
                val: Literal::Empty,
                dst_reg: args_reg,
            },
            pos_call,
        );
        self.code.emit(
            OpCode::Append {
                src_reg: func_reg,
                dst_reg: args_reg,
            },
            pos_call,
        );
        let last_arg = args.pop();
        for arg in args.into_iter() {
            let arg_reg = self.compile_exp(arg)?;
            self.code.emit(
                OpCode::Append {
                    src_reg: arg_reg,
                    dst_reg: args_reg,
                },
                pos_call,
            );
            self.scopes.reg_free(arg_reg);
        }
        if let Some(arg) = last_arg {
            let arg_reg = self.compile_exp(arg)?;
            self.code.emit(
                OpCode::Extend {
                    src_reg: arg_reg,
                    dst_reg: args_reg,
                },
                pos_call,
            );
            self.scopes.reg_free(arg_reg);
        }

        self.code.emit(
            OpCode::Lit {
                val: Literal::String(Rc::new(name.into_bytes())),
                dst_reg: func_reg,
            },
            pos_call,
        );
        self.code.emit(
            OpCode::TableGet {
                tbl_reg: obj_reg,
                ind_reg: func_reg,
                dst_reg: func_reg,
            },
            pos_call,
        );

        self.code.emit(
            OpCode::Call {
                func_reg,
                args_reg,
                ret_reg,
                tail,
            },
            pos_call,
        );

        self.scopes.reg_free(obj_reg);
        self.scopes.reg_free(args_reg);
        self.scopes.reg_free(func_reg);

        Ok(ret_reg)
    }
}