luallaby 0.1.0-alpha.3

use crate::ast::{Exp, PrefixExp, TableField};
use crate::compiler::{offset, Compiler};
use crate::error::Result;
use crate::vm::{BinOp, OpCode, Oper, Value};

impl<'a> Compiler<'a> {
    pub(super) fn compile_exp(&mut self, exp: Exp) -> Result<usize> {
        let oper = self.convert_exp(exp)?;
        Ok(match oper {
            Oper::Field(tbl_reg, ind_reg) => {
                let dst_reg = self.scopes.reg_reserve();
                self.code.emit(OpCode::TableGet {
                    tbl_reg,
                    ind_reg,
                    dst_reg,
                });
                self.scopes.reg_free(tbl_reg);
                self.scopes.reg_free(ind_reg);
                dst_reg
            }
            Oper::Global(name) => {
                let dst_reg = self.scopes.reg_reserve();
                self.code.emit(OpCode::GlobalGet { name, dst_reg });
                dst_reg
            }
            Oper::Local(src_loc) => {
                let dst_reg = self.scopes.reg_reserve();
                self.code.emit(OpCode::LocalGet { src_loc, dst_reg });
                dst_reg
            }
            Oper::Up(src_up) => {
                let dst_reg = self.scopes.reg_reserve();
                self.code.emit(OpCode::UpGet { src_up, dst_reg });
                dst_reg
            }
            Oper::Reg(r) => r,
        })
    }

    fn convert_exp(&mut self, exp: Exp) -> Result<Oper> {
        Ok(match exp {
            Exp::BinOp(lhs, op, rhs) => match op {
                // Short circuit logical operators
                BinOp::And | BinOp::Or => {
                    let lhs_reg = self.compile_exp(*lhs)?;

                    let jump_pos = self.code.pos();
                    self.code.emit(OpCode::Jump { off: 0 }); // Placeholder

                    let rhs_reg = self.compile_exp(*rhs)?;
                    self.code.emit(OpCode::Mov {
                        src_reg: rhs_reg,
                        dst_reg: lhs_reg,
                    });
                    self.scopes.reg_free(rhs_reg);

                    let off = offset(jump_pos, self.code.pos());
                    self.code.set(
                        jump_pos,
                        match op {
                            BinOp::And => OpCode::JumpIfNot {
                                cmp_reg: lhs_reg,
                                off,
                            },
                            BinOp::Or => OpCode::JumpIf {
                                cmp_reg: lhs_reg,
                                off,
                            },
                            _ => unreachable!(),
                        },
                    );

                    Oper::Reg(lhs_reg)
                }
                op => {
                    let lhs = self.convert_exp(*lhs)?;
                    let rhs = self.convert_exp(*rhs)?;
                    let dst_reg = self.scopes.reg_reserve();

                    self.free_oper(&lhs);
                    self.free_oper(&rhs);
                    self.code.emit(OpCode::BinOp {
                        lhs,
                        rhs,
                        op,
                        dst_reg,
                    });

                    Oper::Reg(dst_reg)
                }
            },
            Exp::UnOp(op, exp) => {
                let lhs = self.convert_exp(*exp)?;
                let dst_reg = self.scopes.reg_reserve();

                self.free_oper(&lhs);
                self.code.emit(OpCode::UnOp { lhs, op, dst_reg });

                Oper::Reg(dst_reg)
            }
            Exp::Lit(val) => {
                let dst_reg = self.scopes.reg_reserve();
                self.code.emit(OpCode::Lit { val, dst_reg });
                Oper::Reg(dst_reg)
            }
            Exp::Prefix(exp) => self.compile_prefix_exp_to_oper(exp)?,
            Exp::Table(fields) => {
                let tbl_reg = self.scopes.reg_reserve();
                self.code.emit(OpCode::TableEmpty { dst_reg: tbl_reg });

                let mut index: usize = 1;
                // TODO: Nil and NaN cannot be keys
                let mut fields = fields.into_iter().peekable();
                while let Some(field) = fields.next() {
                    match field {
                        TableField::Index(key, val) => {
                            let ind_reg = self.compile_exp(key)?;
                            let src_reg = self.compile_exp(val)?;

                            self.code.emit(OpCode::TableSet {
                                src_reg,
                                tbl_reg,
                                ind_reg,
                            });

                            self.scopes.reg_free(ind_reg);
                            self.scopes.reg_free(src_reg);
                        }
                        TableField::List(val) => {
                            #[allow(clippy::branches_sharing_code)]
                            if fields.peek().is_none()
                                && matches!(
                                    val,
                                    Exp::Prefix(PrefixExp::Call(..) | PrefixExp::CallMethod(..))
                                        | Exp::Varargs
                                )
                            {
                                // If last element is function call or varargs, extend table
                                let src_reg = self.compile_exp(val)?;

                                self.code.emit(OpCode::TableExtend {
                                    src_reg,
                                    tbl_reg,
                                    ind_from: index,
                                });

                                self.scopes.reg_free(src_reg);
                            } else {
                                let src_reg = self.compile_exp(val)?;
                                let ind_reg = self.scopes.reg_reserve();
                                self.code.emit(OpCode::Lit {
                                    val: Value::int(index as i64),
                                    dst_reg: ind_reg,
                                });
                                index += 1;

                                self.code.emit(OpCode::TableSet {
                                    src_reg,
                                    tbl_reg,
                                    ind_reg,
                                });

                                self.scopes.reg_free(ind_reg);
                                self.scopes.reg_free(src_reg);
                            }
                        }
                    }
                }

                Oper::Reg(tbl_reg)
            }
            Exp::Func(func) => {
                let reg = self.compile_func(func)?;
                Oper::Reg(reg)
            }
            Exp::Varargs => {
                let dst_reg = self.scopes.reg_reserve();
                self.code.emit(OpCode::Varargs { dst_reg });
                Oper::Reg(dst_reg)
            }
        })
    }

    fn free_oper(&mut self, oper: &Oper) {
        match oper {
            Oper::Field(reg_tbl, reg_ind) => {
                self.scopes.reg_free(*reg_tbl);
                self.scopes.reg_free(*reg_ind);
            }
            Oper::Reg(reg) => {
                self.scopes.reg_free(*reg);
            }
            _ => {}
        }
    }
}