luallaby 0.1.0

use crate::ast::{Attr, For, ForGen, ForNum};
use crate::compiler::{offset, Compiler, ScopeType};
use crate::error::Result;
use crate::vm::{BinOp, Literal, OpCode, Oper};

// TODO: Optimize loop runtime
impl Compiler {
    pub(super) fn compile_for(&mut self, r#for: For) -> Result<()> {
        match r#for {
            For::Numeric(f) => self.compile_for_num(f),
            For::Generic(f) => self.compile_for_gen(f),
        }
    }

    fn compile_for_num(&mut self, r#for: ForNum) -> Result<()> {
        self.scope_enter(ScopeType::Loop);

        // Declare control variable and set to init
        let ctrl_reg = self.compile_exp(*r#for.init)?;
        let ctrl = self.declare_local(r#for.name, None, r#for.pos_for);

        // Evaluate limit and step values
        let limit_reg = self.compile_exp(*r#for.limit)?;
        let step_reg = r#for.step.map(|step| self.compile_exp(*step)).transpose()?;

        // Check at start
        let for_start = self.code.current_pc();
        self.code.emit(OpCode::Jump { off: 0 }, r#for.pos_for); // Placeholder

        // Block
        let start = self.code.current_pc();
        self.compile_block(*r#for.block)?;
        self.scopes.mark_loop_end(self.code.current_pc());

        // Check at end
        let for_end = self.code.current_pc();
        self.code.emit(OpCode::Jump { off: 0 }, r#for.pos_for); // Placeholder

        let end = self.code.current_pc();
        self.code.set_op(
            for_start,
            OpCode::ForNum {
                ctrl_reg,
                limit_reg,
                step_reg,
                off_start: offset(for_start, start),
                off_end: offset(for_start, end),
                ctrl_loc: ctrl,
            },
        );
        self.code.set_op(
            for_end,
            OpCode::ForNum {
                ctrl_reg,
                limit_reg,
                step_reg,
                off_start: offset(for_end, start),
                off_end: offset(for_end, end),
                ctrl_loc: ctrl,
            },
        );

        self.scopes.reg_free(ctrl_reg);
        self.scopes.reg_free(limit_reg);
        if let Some(step_reg) = step_reg {
            self.scopes.reg_free(step_reg);
        }

        self.scope_leave(ScopeType::Loop, r#for.pos_end)?;

        Ok(())
    }

    fn compile_for_gen(&mut self, mut r#for: ForGen) -> Result<()> {
        // Evaluate expressions
        let exp_reg = self.scopes.reg_reserve();
        let args_reg = self.scopes.reg_reserve();

        self.code.emit(
            OpCode::Lit {
                val: Literal::Empty,
                dst_reg: exp_reg,
            },
            r#for.pos_for,
        );
        let exp_last = r#for.exps.pop();
        for exp in r#for.exps {
            let src_reg = self.compile_exp(exp)?;
            self.code.emit(
                OpCode::Append {
                    src_reg,
                    dst_reg: exp_reg,
                },
                r#for.pos_for,
            );
            self.scopes.reg_free(src_reg);
        }
        if let Some(exp) = exp_last {
            let src_reg = self.compile_exp(exp)?;
            self.code.emit(
                OpCode::Extend {
                    src_reg,
                    dst_reg: exp_reg,
                },
                r#for.pos_for,
            );
            self.scopes.reg_free(src_reg);
        }

        self.scope_enter(ScopeType::Loop);

        // Declare variables
        let mut vars = Vec::new();
        for name in r#for.names.into_iter() {
            let loc = self.declare_local(name, None, r#for.pos_for);
            vars.push(loc);
        }

        // Initialize variables
        let iter = self.declare_local("(for state)".to_string(), None, r#for.pos_for);
        self.code.emit(
            OpCode::MovMult {
                src_reg: exp_reg,
                ind: 0,
                dst_reg: args_reg,
            },
            r#for.pos_for,
        );
        self.code.emit(
            OpCode::LocalSet {
                src_reg: args_reg,
                dst_loc: iter,
            },
            r#for.pos_for,
        );

        let state = self.declare_local("(for state)".to_string(), None, r#for.pos_for);
        self.code.emit(
            OpCode::MovMult {
                src_reg: exp_reg,
                ind: 1,
                dst_reg: args_reg,
            },
            r#for.pos_for,
        );
        self.code.emit(
            OpCode::LocalSet {
                src_reg: args_reg,
                dst_loc: state,
            },
            r#for.pos_for,
        );

        // Declare empty local to emulate Lua behavior
        self.declare_local("(for state)".to_string(), None, r#for.pos_for);
        let ctrl = vars[0];
        self.code.emit(
            OpCode::MovMult {
                src_reg: exp_reg,
                ind: 2,
                dst_reg: args_reg,
            },
            r#for.pos_for,
        );
        self.code.emit(
            OpCode::LocalSet {
                src_reg: args_reg,
                dst_loc: ctrl,
            },
            r#for.pos_for,
        );

        let close = self.declare_local("(for state)".to_string(), Some(Attr::Close), r#for.pos_for);
        self.code.emit(
            OpCode::MovMult {
                src_reg: exp_reg,
                ind: 3,
                dst_reg: args_reg,
            },
            r#for.pos_for,
        );
        self.code.emit(
            OpCode::LocalSet {
                src_reg: args_reg,
                dst_loc: close,
            },
            r#for.pos_for,
        );

        // Start of loop
        let start = self.code.current_pc();

        // Call iterator
        self.code.emit(
            OpCode::Lit {
                val: Literal::Empty,
                dst_reg: args_reg,
            },
            r#for.pos_call,
        );
        self.code.emit(
            OpCode::LocalGet {
                src_loc: state,
                dst_reg: exp_reg,
            },
            r#for.pos_call,
        );
        self.code.emit(
            OpCode::Append {
                src_reg: exp_reg,
                dst_reg: args_reg,
            },
            r#for.pos_call,
        );
        self.code.emit(
            OpCode::LocalGet {
                src_loc: ctrl,
                dst_reg: exp_reg,
            },
            r#for.pos_call,
        );
        self.code.emit(
            OpCode::Append {
                src_reg: exp_reg,
                dst_reg: args_reg,
            },
            r#for.pos_call,
        );
        self.code.emit(
            OpCode::LocalGet {
                src_loc: iter,
                dst_reg: exp_reg,
            },
            r#for.pos_call,
        );
        self.code.emit(
            OpCode::Call {
                func_reg: exp_reg,
                args_reg,
                ret_reg: exp_reg,
                tail: false,
            },
            r#for.pos_call,
        );

        // Assign variables
        for (ind, var) in vars.into_iter().enumerate() {
            self.code.emit(
                OpCode::MovMult {
                    src_reg: exp_reg,
                    ind,
                    dst_reg: args_reg,
                },
                r#for.pos_for,
            );
            self.code.emit(
                OpCode::LocalSet {
                    src_reg: args_reg,
                    dst_loc: var,
                },
                r#for.pos_for,
            );
        }

        // Check control variable
        self.code.emit(
            OpCode::LocalGet {
                src_loc: ctrl,
                dst_reg: exp_reg,
            },
            r#for.pos_for,
        );
        self.code.emit(
            OpCode::Lit {
                val: Literal::Nil,
                dst_reg: args_reg,
            },
            r#for.pos_for,
        );
        self.code.emit(
            OpCode::BinOp {
                lhs: Oper::Reg(exp_reg),
                rhs: Oper::Reg(args_reg),
                op: BinOp::Eq,
                dst_reg: exp_reg,
            },
            r#for.pos_for,
        );

        // Jump to end if false, placeholder
        let jump_end = self.code.current_pc();
        self.code.emit(OpCode::Jump { off: 0 }, r#for.pos_for); // Placeholder

        // Block
        self.compile_block(*r#for.block)?;
        self.scopes.mark_loop_end(self.code.current_pc());

        // Jump back to start
        let off = offset(self.code.current_pc(), start);
        self.code
            .emit(OpCode::Jump { off }, self.code.get_pos_last());

        // Fill in jump to end
        let off = offset(jump_end, self.code.current_pc());
        self.code.set_op(
            jump_end,
            OpCode::JumpIf {
                cmp_reg: exp_reg,
                off,
            },
        ); // reg_exp at point of insertion contains comparison result

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

        self.scope_leave(ScopeType::Loop, r#for.pos_end)?;

        Ok(())
    }
}