luallaby 0.1.0

use std::collections::HashMap;
use std::mem;
use std::rc::Rc;

mod assignment;
mod block;
mod chunk;
mod expression;
mod r#for;
mod function;
mod r#if;
mod local;
mod prefix_exp;
mod repeat;
mod statement;
mod r#while;

use crate::ast::{Attr, Chunk};
use crate::bitset::BitSet;
use crate::error::{LuaError, Result};
use crate::lexer::Pos;
use crate::vm::{Code, FuncObject, OpCode, VarType};

pub const LUA_ENV: &str = "_ENV";

#[derive(Clone, Copy, Debug)]
pub enum LocalType {
    Local,
    Param,
}

pub struct Compiler {
    // Tracking local variables
    scopes: Scopes,
    code_stack: Vec<Code>,
    code: Code,
    source: Rc<Vec<u8>>, // TODO: Error position in all errors thrown in compiler/parse
}

impl Compiler {
    pub fn compile(chunk: Chunk, source: Rc<Vec<u8>>) -> Result<FuncObject> {
        let mut compiler = Self {
            scopes: Scopes::new(),
            code_stack: Vec::new(),
            code: Code::new(),
            source,
        };

        compiler.compile_chunk(chunk)
    }

    fn scope_enter(&mut self, typ: ScopeType) {
        self.scopes.open(typ);
    }

    fn code_push(&mut self) {
        let push = mem::replace(&mut self.code, Code::new());
        self.code_stack.push(push);
    }

    fn declare_local(&mut self, name: String, attr: Option<Attr>, pos: Pos) -> usize {
        let dst_loc = self.scopes.declare_local(name.clone(), attr);
        self.code.emit(
            OpCode::LocalAlloc {
                dst_loc,
                name,
                to_close: matches!(attr, Some(Attr::Close)),
            },
            pos,
        );
        dst_loc
    }

    fn scope_leave(&mut self, typ: ScopeType, pos: Pos) -> Result<Option<Scope>> {
        let (scope, sub) = self.scopes.close(typ);

        // Replace break placeholders
        for brk in sub.breaks.into_iter() {
            let off = offset(brk.pc, self.code.current_pc());
            self.code.set_op(
                brk.pc,
                OpCode::JumpClose {
                    loc_from: sub.var_count,
                    off,
                },
            );
        }

        // Replace goto placeholders, as far as can be resolved
        let end_of_block = sub.end.unwrap_or_else(|| self.code.current_pc());
        for goto in sub.gotos.into_iter() {
            match sub.labels.get(&goto.label) {
                Some(label) => {
                    // Jumping into the scope of a local variable is not allowed,
                    // unless jumping to the end of a block,
                    // every goto has a 'taint' marker indicating if it is jumping over a new variable declaration,
                    // this taint is set when a new local is declared and the goto does not yet have a label to jump to
                    if let Some(name) = goto.taint {
                        if end_of_block != label.pc {
                            return err!(LuaError::GotoJumpInLocal(goto.label, name));
                        }
                    }
                    let off = offset(goto.pc, label.pc);
                    self.code.set_op(
                        goto.pc,
                        OpCode::JumpClose {
                            loc_from: label.var_count,
                            off,
                        },
                    );
                }
                // Bubble to parent scope, except if function scope was closed
                None => {
                    if scope.is_some() {
                        return err!(LuaError::GotoLabelNotFound(
                            goto.label,
                            self.code.get_pos(goto.pc).line()
                        ));
                    } else {
                        self.scopes.goto_insert(goto.label, goto.pc);
                    }
                }
            };
        }

        // Close variables that went out of scope, except when closing a function scope,
        // when popping a frame off the stack, the VM should take care of closing any locals that need to be closed
        if !matches!(typ, ScopeType::Func) {
            self.code.emit(
                OpCode::LocalClose {
                    loc_from: sub.var_count,
                },
                pos,
            );
        }

        Ok(scope)
    }

    fn code_pop(&mut self) -> Code {
        let pop = self.code_stack.pop().unwrap();
        mem::replace(&mut self.code, pop)
    }
}

enum Variable {
    Direct(VarType, Option<Attr>),
    Global(VarType),
}

struct Scopes {
    scopes: Vec<Scope>,
}

impl Scopes {
    fn new() -> Self {
        Self { scopes: Vec::new() }
    }

    fn open(&mut self, typ: ScopeType) {
        match typ {
            ScopeType::Func => {
                self.scopes.push(if self.scopes.is_empty() {
                    let mut scope = Scope::new();
                    scope.declare_up(
                        LUA_ENV.to_string(),
                        UpDesc {
                            upper: VarType::Up(0),
                            attr: None,
                        },
                    );
                    scope
                } else {
                    Scope::new()
                });
            }
            t => self.scopes.last_mut().unwrap().open(t),
        }
    }

    fn close(&mut self, typ: ScopeType) -> (Option<Scope>, SubScope) {
        let closed_sub = self.scopes.last_mut().unwrap().close();

        assert!(closed_sub.typ == typ);
        let closed_func = match closed_sub.typ {
            ScopeType::Func => {
                let closed_func = self.scopes.pop().unwrap();
                assert!(closed_func.regs.is_empty());
                assert!(closed_func.sub_scopes.is_empty());
                Some(closed_func)
            }
            _ => None,
        };
        (closed_func, closed_sub)
    }

    fn reg_reserve(&mut self) -> usize {
        self.scopes.last_mut().unwrap().regs.reserve()
    }

    fn reg_free(&mut self, ind: usize) {
        self.scopes.last_mut().unwrap().regs.free(ind);
    }

    fn get(&mut self, name: &str) -> Variable {
        match self.get_defined(name) {
            Some((typ, attr)) => Variable::Direct(typ, attr),
            None => Variable::Global(self.get_defined(LUA_ENV).unwrap().0),
        }
    }

    fn get_defined(&mut self, name: &str) -> Option<(VarType, Option<Attr>)> {
        // Find variable
        for (depth, scope) in self.scopes.iter().rev().enumerate() {
            if let Some(var) = scope.get_var(name) {
                let attr = scope.get_attr(var);
                if depth == 0 {
                    // Local found in current scope, just return
                    return Some((*var, attr));
                } else {
                    // Insert upvalues in all scopes above where local was declared
                    let len = self.scopes.len();
                    let mut upper = *var;
                    for scope in self.scopes.iter_mut().skip(len - depth) {
                        let up = scope.declare_up(name.to_string(), UpDesc { upper, attr });
                        upper = VarType::Up(up);
                    }
                    return Some((upper, attr));
                }
            }
        }
        None
    }

    fn declare_params(&mut self, params: Vec<String>) {
        self.scopes.last_mut().unwrap().declare_params(params);
    }

    fn declare_local(&mut self, name: String, attr: Option<Attr>) -> usize {
        self.scopes.last_mut().unwrap().declare_local(LocalDesc {
            name,
            typ: LocalType::Local,
            attr,
        })
    }

    fn has_to_be_closed(&self) -> bool {
        self.scopes
            .last()
            .unwrap()
            .locals
            .iter()
            .any(|loc| matches!(loc.attr, Some(Attr::Close)))
    }

    fn break_insert(&mut self, pc: usize) {
        self.scopes
            .last_mut()
            .unwrap()
            .break_insert(BreakDesc { pc });
    }

    fn goto_insert(&mut self, label: String, pc: usize) {
        self.scopes.last_mut().unwrap().goto_insert(GotoDesc {
            label,
            pc,
            taint: None,
        });
    }

    fn label_insert(&mut self, name: String, pc: usize, pos: Pos) -> Result<()> {
        self.scopes.last_mut().unwrap().label_insert(name, pc, pos)
    }

    fn mark_loop_end(&mut self, pc: usize) {
        self.scopes.last_mut().unwrap().mark_loop_end(pc);
    }
}

pub struct Scope {
    sub_scopes: Vec<SubScope>,
    // Registers
    regs: BitSet,
    // Variable counts
    locals_cap: usize,
    // Variable slots
    locals: Vec<LocalDesc>,
    ups: Vec<UpDesc>,
}

struct SubScope {
    typ: ScopeType,
    // How many variables were declared at point of scope enter, indicating how many should be closed on leaving
    var_count: usize,
    // Variable search dict
    dict: HashMap<String, VarType>,
    // Breaks
    breaks: Vec<BreakDesc>,
    // Gotos
    gotos: Vec<GotoDesc>,
    labels: HashMap<String, LabelDesc>,
    // End of block position, used for allowing gotos jumping over local declarations to end of loop block
    end: Option<usize>,
}

#[derive(Clone, Copy, PartialEq)]
enum ScopeType {
    Func,
    Loop,
    Do,
}

struct LocalDesc {
    name: String,
    typ: LocalType,
    attr: Option<Attr>,
}

struct UpDesc {
    upper: VarType,
    attr: Option<Attr>,
}

struct BreakDesc {
    pc: usize,
}

struct GotoDesc {
    label: String,
    pc: usize,
    taint: Option<String>,
}

struct LabelDesc {
    pc: usize,
    var_count: usize,
    pos: Pos,
}

impl Scope {
    fn new() -> Self {
        Self {
            regs: BitSet::new(),
            sub_scopes: vec![SubScope {
                typ: ScopeType::Func,
                var_count: 0,
                dict: HashMap::new(),
                breaks: Vec::new(),
                gotos: Vec::new(),
                labels: HashMap::new(),
                end: None,
            }],
            locals_cap: 0,
            locals: Vec::new(),
            ups: Vec::new(),
        }
    }

    fn open(&mut self, typ: ScopeType) {
        assert!(!matches!(typ, ScopeType::Func));
        self.sub_scopes.push(SubScope {
            typ,
            var_count: self.locals.len(),
            dict: HashMap::new(),
            breaks: Vec::new(),
            gotos: Vec::new(),
            labels: HashMap::new(),
            end: None,
        });
    }

    fn close(&mut self) -> SubScope {
        let sub = self.sub_scopes.pop().unwrap();
        // Remove variables that went out of scope
        self.locals.truncate(sub.var_count);
        sub
    }

    fn get_var(&self, name: &str) -> Option<&VarType> {
        for sub in self.sub_scopes.iter().rev() {
            if let Some(var) = sub.dict.get(name) {
                return Some(var);
            }
        }
        None
    }

    fn get_attr(&self, var: &VarType) -> Option<Attr> {
        match var {
            VarType::Local(ind) => self.locals[*ind].attr,
            VarType::Up(ind) => self.ups[*ind].attr,
        }
    }

    fn declare_params(&mut self, params: Vec<String>) {
        let sub = self.sub_scopes.last_mut().unwrap();
        assert!(self.locals.is_empty() && self.locals_cap == 0 && sub.var_count == 0);

        for name in params.into_iter() {
            let ind = self.locals.len();
            self.locals.push(LocalDesc {
                name: name.clone(),
                typ: LocalType::Param,
                attr: None,
            });
            sub.dict.insert(name, VarType::Local(ind));
        }

        self.locals_cap = self.locals.len();
        // Compensate for parameters, that don't need to be closed
        sub.var_count = self.locals.len();
    }

    fn declare_local(&mut self, local: LocalDesc) -> usize {
        assert!(matches!(local.typ, LocalType::Local));
        let ind = self.locals.len();

        let sub = self.sub_scopes.last_mut().unwrap();
        // When a goto does not have a label to jump to yet, and it hasn't been tained, taint it
        for goto in sub.gotos.iter_mut() {
            // TODO: This check does not need to happen when gotos already have a matching label, maybe collect those
            if !sub.labels.contains_key(&goto.label) {
                goto.taint.get_or_insert(local.name.clone());
            }
        }
        sub.dict.insert(local.name.clone(), VarType::Local(ind));

        self.locals.push(local);
        self.locals_cap = self.locals_cap.max(self.locals.len());
        ind
    }

    fn declare_up(&mut self, name: String, up: UpDesc) -> usize {
        let ind = self.ups.len();
        self.ups.push(up);
        // Always insert upvalue on bottom of sub-scope stack,
        // the same name will always resolve to that upvalue,
        // regardless of how many sub-scopes we have in this function,
        // if we would insert on top of stack then when that sub-scope
        // is dropped we would have to resolve and insert the same upvalue again
        self.sub_scopes
            .first_mut()
            .unwrap()
            .dict
            .insert(name, VarType::Up(ind));
        ind
    }

    fn break_insert(&mut self, brk: BreakDesc) {
        for sub in self.sub_scopes.iter_mut().rev() {
            if let ScopeType::Loop = sub.typ {
                sub.breaks.push(brk);
                return;
            }
        }
        panic!("not in loop")
    }

    fn goto_insert(&mut self, goto: GotoDesc) {
        self.sub_scopes.last_mut().unwrap().gotos.push(goto);
    }

    fn label_insert(&mut self, name: String, pc: usize, pos: Pos) -> Result<()> {
        for sub in self.sub_scopes.iter().rev() {
            if let Some(label) = sub.labels.get(&name) {
                return err!(LuaError::GotoLabelRepeat(name, label.pos.line()));
            }
        }
        let sub = self.sub_scopes.last_mut().unwrap();
        sub.labels.insert(
            name,
            LabelDesc {
                pc,
                var_count: self.locals.len(),
                pos,
            },
        );
        Ok(())
    }

    fn mark_loop_end(&mut self, pc: usize) {
        let sub = self.sub_scopes.last_mut().unwrap();
        assert!(matches!(sub.typ, ScopeType::Loop));
        sub.end = Some(pc);
    }
}

/// Expects position of jump and position to jump to
fn offset(jump: usize, target: usize) -> i64 {
    // Compiler::pos() gets the position of the next emitted opcode,
    // usually one would save the result from pos(), emit a placeholder,
    // and then later use the saved pos to fill placeholder position,
    // this means that we need to adjust by one here
    target as i64 - jump as i64 - 1
}