compiler.rs

use std::{
    collections::{HashMap, HashSet},
    iter,
};

use cstree::{self, interning::TokenInterner};

use super::{
    super::parser::syntax::{Root, Stmt},
    bytecode::{Capture, Instruction},
    error::Error,
    gc::Heap,
    value::Value,
    vm::Chunk,
    ObjectPtr,
    ScriptData,
};
use crate::{
    inst,
    parser::syntax::{
        Assign,
        BinaryOp,
        BinaryOperator,
        Break,
        Decl,
        DeclModifier,
        Do,
        Expr,
        ForGen,
        ForNum,
        Func,
        FuncCall,
        FuncExpr,
        Goto,
        Ident,
        If,
        Index,
        Label,
        Literal,
        LiteralValue,
        MethodCall,
        PrefixOp,
        PrefixOperator,
        Repeat,
        Return,
        Table,
        TableEntry,
        While,
    },
};

macro_rules! try_err {
    ($e:expr, $err:expr) => {
        match $e {
            Some(v) => v,
            None => return Err(Error::new($err)),
        }
    };
}

struct VariableData {
    location: u8,
    is_const: bool,
}

struct Ctx<'a> {
    interner: &'a TokenInterner,
    chunk_alloc: &'a mut dyn FnMut(Chunk) -> u32,
    string_alloc: &'a mut dyn FnMut(&[u8]) -> ObjectPtr,
    chunk: Chunk,
    control_end_label: Vec<u16>,
    return_label: u16,
    scope: Vec<HashMap<String, VariableData>>,
    scope_destructors: Vec<HashSet<u8>>,
    labels: HashMap<String, u16>,
    capture: &'a mut dyn FnMut(&str) -> Option<u8>,
    captured: Vec<u8>,
}

impl<'a> Ctx<'a> {
    fn emit(&mut self, instruction: Instruction) {
        self.chunk.tape.push(instruction);
    }

    fn alloc_register(&mut self) -> u8 {
        let reg = self.chunk.register_count as u8;
        self.chunk.register_count += 1;
        reg
    }

    fn next_offset(&self) -> usize {
        self.chunk.tape.len()
    }

    fn new_label(&mut self) -> u16 {
        let label = self.chunk.labels.len();
        self.chunk.labels.push(0);
        label as u16
    }

    fn set_label(&mut self, label: u16, offset: usize) {
        self.chunk.labels[label as usize] = offset;
    }

    fn set_control_end_label(&mut self, label: u16) {
        self.control_end_label.push(label);
    }

    fn get_control_end_label(&mut self) -> Option<u16> {
        self.control_end_label.last().copied()
    }

    fn clear_control_end_label(&mut self) {
        self.control_end_label.pop();
    }

    fn push_scope(&mut self) {
        self.scope.push(HashMap::new());
        self.scope_destructors.push(HashSet::new());
    }

    fn pop_scope(&mut self) -> Result<HashSet<u8>, Error> {
        try_err!(self.scope.pop(), "missing scope");
        Ok(try_err!(self.scope_destructors.pop(), "missing scope"))
    }

    fn do_close(&mut self, register: u8) -> Result<(), Error> {
        try_err!(self.scope_destructors.last_mut(), "missing optional").insert(register);
        Ok(())
    }

    fn define(&mut self, name: String, data: VariableData) -> Result<(), Error> {
        let scope = try_err!(self.scope.last_mut(), "missing optional");
        scope.insert(name, data);
        Ok(())
    }

    fn resolve(&mut self, name: &str) -> Option<&VariableData> {
        for scope in self.scope.iter().rev() {
            if let Some(data) = scope.get(name) {
                return Some(data);
            }
        }

        None
    }

    fn alloc_constant(&mut self, constant: Value) -> u16 {
        let index = self.chunk.constants.len();
        self.chunk.constants.push(constant);
        index as u16
    }

    fn add_goto(&mut self, name: String, label: u16) {
        self.labels.insert(name, label);
    }

    fn resolve_goto(&self, name: &str) -> Option<u16> {
        self.labels.get(name).copied()
    }
}

fn get_capture(
    scope: &[HashMap<String, VariableData>],
    capture: &mut dyn FnMut(&str) -> Option<u8>,
    captured: &mut HashMap<String, (Capture, u8)>,
    name: &str,
) -> Option<u8> {
    if let Some((_, id)) = captured.get(name) {
        return Some(*id);
    }

    for scope in scope.iter().rev() {
        if let Some(data) = scope.get(name) {
            let id = captured.len() as u8;
            captured.insert(
                name.to_owned(),
                (
                    Capture {
                        location: data.location,
                        is_local: true,
                    },
                    id,
                ),
            );

            return Some(id);
        }
    }

    if let Some(idx) = capture(name) {
        let id = captured.len() as u8;
        captured.insert(
            name.to_owned(),
            (
                Capture {
                    location: idx,
                    is_local: false,
                },
                id,
            ),
        );

        return Some(id);
    }

    None
}

fn scope_lexical<F>(ctx: &mut Ctx, compile: F) -> Result<(), Error>
where
    F: FnOnce(&mut Ctx) -> Result<(), Error>,
{
    ctx.push_scope();
    compile(ctx)?;

    for register in ctx.pop_scope()? {
        ctx.emit(inst![close => register]);
    }

    Ok(())
}

fn scope_break<F>(ctx: &mut Ctx, compile: F) -> Result<(), Error>
where
    F: FnOnce(&mut Ctx) -> Result<(), Error>,
{
    let label_control_end = ctx.new_label();
    ctx.set_control_end_label(label_control_end);
    compile(ctx)?;
    ctx.set_label(label_control_end, ctx.next_offset());
    ctx.clear_control_end_label();
    Ok(())
}

fn scope_lexical_break<F>(ctx: &mut Ctx, compile: F) -> Result<(), Error>
where
    F: FnOnce(&mut Ctx) -> Result<(), Error>,
{
    scope_lexical(ctx, |ctx| scope_break(ctx, compile))
}

pub fn compile(
    tree: Root,
    interner: &TokenInterner,
    id: u32,
    env: ObjectPtr,
    heap: &Heap,
    chunk_alloc: &mut dyn FnMut(Chunk) -> u32,
    string_alloc: &mut dyn FnMut(&[u8]) -> ObjectPtr,
) -> Result<ScriptData, Error> {
    let mut chunks = Vec::new();
    let chunk_alloc = &mut |chunk: Chunk| {
        let id = chunk_alloc(chunk);
        chunks.push(id);
        id
    };

    let chunk = compile_chunk(
        interner,
        chunk_alloc,
        string_alloc,
        &mut tree.block(),
        &mut iter::empty(),
        &mut |_| None,
        "__main".to_owned(),
    )?;

    let root = heap.allocate_function(chunk, 0, &mut iter::empty());
    Ok(ScriptData {
        id,
        root,
        env,
        chunks,
    })
}

fn compile_chunk(
    interner: &TokenInterner,
    chunk_alloc: &mut dyn FnMut(Chunk) -> u32,
    string_alloc: &mut dyn FnMut(&[u8]) -> ObjectPtr,
    stmts: &mut dyn Iterator<Item = Stmt>,
    args: &mut dyn Iterator<Item = Result<String, Error>>,
    capture: &mut dyn FnMut(&str) -> Option<u8>,
    name: String,
) -> Result<u32, Error> {
    let chunk = Chunk {
        tape: Vec::new(),
        constants: Vec::new(),
        register_count: 0,
        arity: 0,
        labels: Vec::new(),
        name,
    };

    let mut ctx = Ctx {
        interner,
        chunk_alloc,
        string_alloc,
        chunk,
        control_end_label: Vec::new(),
        return_label: 0,
        scope: Vec::new(),
        scope_destructors: Vec::new(),
        labels: HashMap::new(),
        capture,
        captured: Vec::new(),
    };

    scope_lexical(&mut ctx, |ctx| {
        for arg in args {
            let register = ctx.alloc_register();
            ctx.emit(inst!(apop => register));
            ctx.define(
                arg?,
                VariableData {
                    location: register,
                    is_const: false,
                },
            )?;
        }

        ctx.emit(inst![aclear]);
        ctx.return_label = ctx.new_label();
        for item in stmts {
            compile_stmt(ctx, item)?;
        }

        ctx.set_label(ctx.return_label, ctx.next_offset());
        let captured = ctx.captured.clone();
        for register in captured {
            ctx.emit(inst![uclo => register]);
        }
        ctx.emit(inst![ret]);
        Ok(())
    })?;

    let Ctx {
        chunk_alloc, chunk, ..
    } = ctx;

    Ok(chunk_alloc(chunk))
}

fn compile_stmt(ctx: &mut Ctx, item: Stmt) -> Result<(), Error> {
    match item {
        Stmt::Label(item) => compile_label(ctx, item),
        Stmt::Goto(item) => compile_goto(ctx, item),
        Stmt::Decl(item) => compile_decl(ctx, item),
        Stmt::Assign(item) => compile_assign(ctx, item),
        Stmt::Func(item) => compile_func(ctx, item),
        Stmt::Expr(item) => compile_expr(ctx, item).map(|_| ()),
        Stmt::Break(item) => compile_break(ctx, item),
        Stmt::Return(item) => compile_return(ctx, item),
        Stmt::Do(item) => compile_do(ctx, item),
        Stmt::While(item) => compile_while(ctx, item),
        Stmt::Repeat(item) => compile_repeat(ctx, item),
        Stmt::If(item) => compile_if(ctx, item),
        Stmt::ForNum(item) => compile_for_num(ctx, item),
        Stmt::ForGen(item) => compile_for_gen(ctx, item),
    }
}

fn compile_label(ctx: &mut Ctx, item: Label) -> Result<(), Error> {
    let name = try_err!(item.name().unwrap().name(ctx.interner), "missing optional");
    let label = ctx.new_label();
    ctx.set_label(label, ctx.next_offset());
    ctx.add_goto(name.to_string(), label);
    Ok(())
}

fn compile_goto(ctx: &mut Ctx, item: Goto) -> Result<(), Error> {
    let name = try_err!(item.label().unwrap().name(ctx.interner), "missing optional");
    let label = try_err!(ctx.resolve_goto(name), "undefined label");
    ctx.emit(inst![jmp => label]);
    Ok(())
}

fn compile_decl(ctx: &mut Ctx, item: Decl) -> Result<(), Error> {
    if let Some(func) = item.function() {
        let name = if let Some(Expr::Ident(ident)) = func.target() {
            try_err!(ident.name(ctx.interner), "missing optional")
        } else {
            unreachable!();
        };

        let register = ctx.alloc_register();
        ctx.define(
            name.to_string(),
            VariableData {
                location: register,
                is_const: false,
            },
        )?;

        compile_func(ctx, func)?;
        return Ok(());
    }

    let want = try_err!(item.targets(), "missing optional").count();
    let values = &mut try_err!(item.values(), "missing optional");
    let values_len = try_err!(item.values(), "missing optional").count();
    let sources = prepare_assign_sources(ctx, values, values_len, want)?;

    for (target, value) in try_err!(item.targets(), "missing optional").zip(sources) {
        let name = try_err!(
            target.name().unwrap().name(ctx.interner),
            "missing optional"
        );
        let modifier = target.modifier();
        let register = ctx.alloc_register();
        ctx.emit(inst![mov => register, value]);

        if matches!(modifier, Some(DeclModifier::Close)) {
            ctx.do_close(register)?;
        }

        ctx.define(
            name.to_string(),
            VariableData {
                location: register,
                is_const: matches!(modifier, Some(DeclModifier::Const)),
            },
        )?;
    }

    Ok(())
}

fn compile_assign(ctx: &mut Ctx, item: Assign) -> Result<(), Error> {
    let targets = try_err!(item.targets(), "missing optional");
    let values = &mut try_err!(item.values(), "missing optional");
    let values_len = try_err!(item.values(), "missing optional").count();
    let want = try_err!(item.targets(), "missing optional").count();
    let sources = prepare_assign_sources(ctx, values, values_len, want)?;

    for (target, value) in targets.zip(sources) {
        compile_assign_lhs_expr(ctx, target, value)?;
    }

    Ok(())
}

fn prepare_assign_sources(
    ctx: &mut Ctx,
    values: &mut dyn Iterator<Item = Expr>,
    values_len: usize,
    mut want: usize,
) -> Result<Vec<u8>, Error> {
    let mut sources = Vec::new();
    for (i, expr) in values.enumerate() {
        if let Expr::FuncCall(call) = expr {
            if values_len == i + 1 {
                for register in compile_expr_func_call(ctx, call, want)? {
                    sources.push(register);
                }
            } else {
                sources.push(compile_expr_func_call(ctx, call, 1)?[0]);
                want -= 1;
            }
        } else {
            let register = compile_expr(ctx, expr)?;
            sources.push(register);
            want -= 1;
        }
    }

    Ok(sources)
}

fn compile_assign_lhs_expr(ctx: &mut Ctx, target: Expr, value: u8) -> Result<(), Error> {
    let assign_table = |ctx: &mut Ctx, object, index| -> Result<(), Error> {
        let object = compile_expr(ctx, object)?;
        let index = compile_expr(ctx, index)?;
        ctx.emit(inst![tset => object, index, value]);
        Ok(())
    };

    let not_valid = || Err(Error::new("expression target not valid"));

    match target {
        Expr::Ident(item) => {
            let name = try_err!(item.name(ctx.interner), "missing optional");
            if let Some(data) = ctx.resolve(name) {
                if data.is_const {
                    return Err(Error::new("cannot reassign constant variable"));
                }

                let register = data.location;
                ctx.emit(inst![mov => register, value]);
                return Ok(());
            }

            if let Some(upvalue) = (ctx.capture)(name) {
                ctx.emit(inst![uset => value, upvalue]);
                return Ok(());
            }

            let index = compile_ident_field(ctx, item)?;
            ctx.emit(inst![gset => index, value]);
        },
        Expr::Index(item) => {
            let object = try_err!(item.target(), "missing optional");
            let index = try_err!(item.index(), "missing optional");
            assign_table(ctx, object, index)?;
        },
        Expr::BinaryOp(item) => {
            if !matches!(
                try_err!(item.op(), "missing optional"),
                BinaryOperator::Property | BinaryOperator::Method
            ) {
                return not_valid();
            }

            let object = try_err!(item.lhs(), "missing optional");
            let index = try_err!(item.rhs(), "missing optional");
            assign_table(ctx, object, index)?;
        },
        _ => return not_valid(),
    }

    Ok(())
}

fn compile_func(ctx: &mut Ctx, item: Func) -> Result<(), Error> {
    let name = try_err!(item.name(ctx.interner), "missing optional");
    let target = try_err!(item.target(), "missing optional");
    let is_method = if let Expr::BinaryOp(item) = &target {
        item.op() == Some(BinaryOperator::Method)
    } else {
        false
    };

    let mut stmts = try_err!(item.block(), "missing optional");
    let mut args: Vec<String> = item
        .args()
        .unwrap()
        .map(|arg| Ok(try_err!(arg.name(ctx.interner), "missing optional").to_string()))
        .collect::<Result<Vec<String>, Error>>()?;

    if is_method {
        args.insert(0, "self".to_string());
    }

    let mut captured = HashMap::new();
    let chunk = compile_chunk(
        ctx.interner,
        ctx.chunk_alloc,
        ctx.string_alloc,
        &mut stmts,
        &mut args.into_iter().map(Ok),
        &mut |name| get_capture(&ctx.scope, ctx.capture, &mut captured, name),
        name.to_string(),
    )?;

    let capture_list: Vec<Capture> = captured
        .values()
        .map(|(capture, _)| capture)
        .copied()
        .collect();
    for capture in &capture_list {
        if capture.is_local {
            ctx.captured.push(capture.location);
        }
    }

    let register = ctx.alloc_register();
    ctx.emit(inst![fnew => register, chunk as u16, capture_list]);
    compile_assign_lhs_expr(ctx, target, register)?;
    Ok(())
}

fn compile_expr(ctx: &mut Ctx, item: Expr) -> Result<u8, Error> {
    let first = |x: Vec<u8>| x[0];

    match item {
        Expr::Method(item) => compile_expr_method_call(ctx, item, 1).map(first),
        Expr::Ident(item) => compile_expr_ident(ctx, item),
        Expr::Literal(item) => compile_expr_literal(ctx, item),
        Expr::Func(item) => compile_expr_func(ctx, item),
        Expr::Table(item) => compile_expr_table(ctx, item),
        Expr::PrefixOp(item) => compile_expr_prefix_op(ctx, item),
        Expr::BinaryOp(item) => compile_expr_binary_op(ctx, item),
        Expr::FuncCall(item) => compile_expr_func_call(ctx, item, 1).map(first),
        Expr::Index(item) => compile_expr_index(ctx, item),
        Expr::VarArg => unimplemented!(),
    }
}

fn compile_expr_ident(ctx: &mut Ctx, item: Ident) -> Result<u8, Error> {
    let name = try_err!(item.name(ctx.interner), "missing optional");

    if let Some(data) = ctx.resolve(name) {
        return Ok(data.location);
    }

    let register = ctx.alloc_register();
    if let Some(upvalue) = (ctx.capture)(name) {
        ctx.emit(inst![uget => register, upvalue]);
        return Ok(register);
    }

    let idx = compile_ident_field(ctx, item)?;
    ctx.emit(inst![gget => register, idx]);
    Ok(register)
}

fn compile_expr_literal(ctx: &mut Ctx, item: Literal) -> Result<u8, Error> {
    let value = match try_err!(item.value(ctx.interner), "missing optional") {
        LiteralValue::Nil => Value::nil(),
        LiteralValue::Bool(value) => Value::from_bool(value),
        LiteralValue::Int(value) => Value::from_int(value),
        LiteralValue::Float(value) => Value::from_float(value),
        LiteralValue::String(value) => {
            let ptr = (ctx.string_alloc)(&value);
            Value::from_string(ptr)
        },
    };

    let register = ctx.alloc_register();
    let idx = ctx.alloc_constant(value);
    ctx.emit(inst![cload => register, idx]);
    Ok(register)
}

fn compile_expr_func(ctx: &mut Ctx, item: FuncExpr) -> Result<u8, Error> {
    let mut stmts = try_err!(item.block(), "missing optional");
    let mut args = item
        .args()
        .unwrap()
        .map(|arg| Ok(try_err!(arg.name(ctx.interner), "missing optional").to_string()));

    let mut captures = HashMap::new();
    let chunk = compile_chunk(
        ctx.interner,
        ctx.chunk_alloc,
        ctx.string_alloc,
        &mut stmts,
        &mut args,
        &mut |name| get_capture(&ctx.scope, ctx.capture, &mut captures, name),
        "__anon".to_string(),
    )?;

    let capture_list: Vec<Capture> = captures
        .values()
        .map(|(capture, _)| capture)
        .copied()
        .collect();
    let register = ctx.alloc_register();
    ctx.emit(inst![fnew => register, chunk as u16, capture_list]);
    Ok(register)
}

fn compile_ident_field(ctx: &mut Ctx, item: Ident) -> Result<u8, Error> {
    let name = try_err!(item.name(ctx.interner), "missing optional");
    let ptr = (ctx.string_alloc)(name.as_bytes());
    let idx = ctx.alloc_constant(Value::from_string(ptr));
    let idx_register = ctx.alloc_register();
    ctx.emit(inst![cload => idx_register, idx]);
    Ok(idx_register)
}

fn compile_expr_table(ctx: &mut Ctx, item: Table) -> Result<u8, Error> {
    let table = ctx.alloc_register();
    ctx.emit(inst![tnew => table]);
    let mut i = 1;

    for entry in item.entries() {
        match entry {
            TableEntry::Array(item) => {
                let idx = ctx.alloc_constant(Value::from_int(i));
                let key = ctx.alloc_register();
                let value = try_err!(item.value(), "missing optional");
                let value = compile_expr(ctx, value)?;
                ctx.emit(inst![cload => key, idx]);
                ctx.emit(inst![tset => table, key, value]);
                i += 1;
            },
            TableEntry::Map(item) => {
                let field = try_err!(item.field(), "missing optional");
                let key = compile_ident_field(ctx, field)?;
                let value = try_err!(item.value(), "missing optional");
                let value = compile_expr(ctx, value)?;
                ctx.emit(inst![tset => table, key, value]);
            },
            TableEntry::Generic(item) => {
                let key = try_err!(item.index(), "missing optional");
                let key = compile_expr(ctx, key)?;
                let value = try_err!(item.value(), "missing optional");
                let value = compile_expr(ctx, value)?;
                ctx.emit(inst![tset => table, key, value]);
            },
        }
    }

    Ok(table)
}

fn compile_expr_prefix_op(ctx: &mut Ctx, item: PrefixOp) -> Result<u8, Error> {
    let dst = ctx.alloc_register();
    let rhs_expr = try_err!(item.rhs(), "missing optional");
    let rhs = compile_expr(ctx, rhs_expr)?;

    ctx.emit(match try_err!(item.op(), "missing optional") {
        PrefixOperator::None => inst![mov => dst, rhs],
        PrefixOperator::Neg => inst![neg => dst, rhs],
        PrefixOperator::Not => inst![not => dst, rhs],
        PrefixOperator::Len => inst![len => dst, rhs],
        PrefixOperator::BitNot => inst![bitnot => dst, rhs],
    });

    Ok(dst)
}

fn compile_expr_binary_op(ctx: &mut Ctx, item: BinaryOp) -> Result<u8, Error> {
    let dst = ctx.alloc_register();
    let lhs_expr = try_err!(item.lhs(), "missing optional");
    let rhs_expr = try_err!(item.rhs(), "missing optional");
    let lhs = compile_expr(ctx, lhs_expr)?;
    let rhs = compile_expr(ctx, rhs_expr)?;

    let inst = match try_err!(item.op(), "missing optional") {
        BinaryOperator::And => inst![and => dst, lhs, rhs],
        BinaryOperator::Or => inst![or => dst, lhs, rhs],
        BinaryOperator::Add => inst![add => dst, lhs, rhs],
        BinaryOperator::Sub => inst![sub => dst, lhs, rhs],
        BinaryOperator::Mul => inst![mul => dst, lhs, rhs],
        BinaryOperator::Div => inst![div => dst, lhs, rhs],
        BinaryOperator::IntDiv => inst![intdiv => dst, lhs, rhs],
        BinaryOperator::Exp => inst![exp => dst, lhs, rhs],
        BinaryOperator::Mod => inst![mod => dst, lhs, rhs],
        BinaryOperator::BitAnd => inst![bitand => dst, lhs, rhs],
        BinaryOperator::BitOr => inst![bitor => dst, lhs, rhs],
        BinaryOperator::LShift => inst![leftshift => dst, lhs, rhs],
        BinaryOperator::RShift => inst![rightshift => dst, lhs, rhs],
        BinaryOperator::Eq => inst![eq => dst, lhs, rhs],
        BinaryOperator::BitXor => inst![bitxor => dst, lhs, rhs],
        BinaryOperator::NEq => inst![neq => dst, lhs, rhs],
        BinaryOperator::LEq => inst![leq => dst, lhs, rhs],
        BinaryOperator::GEq => inst![leq => dst, rhs, lhs],
        BinaryOperator::Gt => inst![lt => dst, rhs, lhs],
        BinaryOperator::Lt => inst![lt => dst, lhs, rhs],
        BinaryOperator::Property => inst![tget => lhs, dst, rhs],
        BinaryOperator::Method => panic!(),
        BinaryOperator::Concat => inst![concat => dst, lhs, rhs],
    };

    ctx.emit(inst);
    Ok(dst)
}

fn compile_expr_func_call(ctx: &mut Ctx, item: FuncCall, want: usize) -> Result<Vec<u8>, Error> {
    let func = try_err!(item.target(), "missing optional");
    let callee = compile_expr(ctx, func)?;

    #[allow(clippy::needless_collect)]
    let args: Vec<u8> = item
        .args()
        .unwrap()
        .map(|arg| compile_expr(ctx, arg))
        .collect::<Result<_, _>>()?;

    Ok(compile_call(ctx, callee, &mut args.into_iter(), want))
}

fn compile_expr_method_call(
    ctx: &mut Ctx,
    item: MethodCall,
    want: usize,
) -> Result<Vec<u8>, Error> {
    let object = try_err!(item.object(), "missing optional");
    let object = compile_expr(ctx, object)?;
    let method = try_err!(item.method(), "missing optional");
    let idx = compile_ident_field(ctx, method)?;
    let func_register = ctx.alloc_register();
    ctx.emit(inst![tget => object, func_register, idx]);

    #[allow(clippy::needless_collect)]
    let mut args: Vec<u8> = item
        .args()
        .unwrap()
        .map(|arg| compile_expr(ctx, arg))
        .collect::<Result<_, _>>()?;

    args.insert(0, object);
    Ok(compile_call(
        ctx,
        func_register,
        &mut args.into_iter(),
        want,
    ))
}

fn compile_expr_index(ctx: &mut Ctx, item: Index) -> Result<u8, Error> {
    let object = try_err!(item.target(), "missing optional");
    let object = compile_expr(ctx, object)?;
    let index = try_err!(item.index(), "missing optional");
    let index = compile_expr(ctx, index)?;
    let dst = ctx.alloc_register();
    ctx.emit(inst![tget => object, dst, index]);
    Ok(dst)
}

fn compile_call(
    ctx: &mut Ctx,
    callee: u8,
    args: &mut dyn Iterator<Item = u8>,
    want: usize,
) -> Vec<u8> {
    for arg in args {
        ctx.emit(inst![apush => arg]);
    }

    ctx.emit(inst![call => callee]);
    let mut registers = Vec::new();

    for _ in 0..want {
        let register = ctx.alloc_register();
        ctx.emit(inst![apop => register]);
        registers.push(register);
    }

    ctx.emit(inst![aclear]);
    registers
}

fn compile_break(ctx: &mut Ctx, _item: Break) -> Result<(), Error> {
    let label_control_end = try_err!(ctx.get_control_end_label(), "cannot break outside of block");
    if label_control_end == 0 {
        panic!();
    }

    ctx.emit(inst![jmp => label_control_end]);
    Ok(())
}

fn compile_return(ctx: &mut Ctx, item: Return) -> Result<(), Error> {
    for item in try_err!(item.exprs(), "missing optional") {
        let value = compile_expr(ctx, item)?;
        ctx.emit(inst![apush => value]);
    }

    ctx.emit(inst![jmp => ctx.return_label]);
    Ok(())
}

fn compile_do(ctx: &mut Ctx, item: Do) -> Result<(), Error> {
    scope_lexical(ctx, |ctx| {
        for item in try_err!(item.stmts(), "missing optional") {
            compile_stmt(ctx, item)?;
        }

        Ok(())
    })
}

fn compile_while(ctx: &mut Ctx, item: While) -> Result<(), Error> {
    scope_lexical_break(ctx, |ctx| {
        let label_cond = ctx.new_label();
        ctx.emit(inst![jmp => label_cond]);

        let label_body = ctx.new_label();
        ctx.set_label(label_body, ctx.next_offset());

        for item in try_err!(item.block(), "missing optional") {
            compile_stmt(ctx, item)?;
        }

        ctx.set_label(label_cond, ctx.next_offset());
        let cond = compile_expr(ctx, try_err!(item.cond(), "missing optional"))?;

        ctx.emit(inst![ist => cond]);
        ctx.emit(inst![cjmp => label_body]);
        Ok(())
    })
}

fn compile_repeat(ctx: &mut Ctx, item: Repeat) -> Result<(), Error> {
    scope_break(ctx, |ctx| {
        let label_body = ctx.new_label();
        ctx.set_label(label_body, ctx.next_offset());

        scope_lexical(ctx, |ctx| {
            for item in try_err!(item.block(), "missing optional") {
                compile_stmt(ctx, item)?;
            }

            Ok(())
        })?;

        if let Some(cond) = item.cond() {
            let cond = compile_expr(ctx, cond)?;

            ctx.emit(inst![ist => cond]);
            ctx.emit(inst![cjmp => label_body]);
        } else {
            ctx.emit(inst![jmp => label_body]);
        }

        Ok(())
    })
}

fn compile_if(ctx: &mut Ctx, item: If) -> Result<(), Error> {
    scope_lexical(ctx, |ctx| {
        let label_body_end = ctx.new_label();
        let cond = compile_expr(ctx, try_err!(item.cond(), "missing optional"))?;
        let chain = item.else_chain();

        ctx.emit(inst![isf => cond]);
        ctx.emit(inst![cjmp => label_body_end]);

        for item in try_err!(item.stmts(), "missing optional") {
            compile_stmt(ctx, item)?;
        }

        let label_else_end = ctx.new_label();
        if chain.is_some() {
            ctx.emit(inst![jmp => label_else_end]);
        }

        ctx.set_label(label_body_end, ctx.next_offset());
        if let Some(chain) = chain {
            if let Some(stmts) = chain.else_block() {
                for item in stmts {
                    compile_stmt(ctx, item)?;
                }
            } else if let Some(item) = chain.elseif_block() {
                compile_if(ctx, item)?;
            }
        }

        ctx.set_label(label_else_end, ctx.next_offset());
        Ok(())
    })
}

fn compile_for_num(ctx: &mut Ctx, item: ForNum) -> Result<(), Error> {
    scope_lexical_break(ctx, |ctx| {
        let (counter_name, initial) = try_err!(item.counter(), "missing optional");
        let counter_name =
            try_err!(counter_name.name(ctx.interner), "missing optional").to_string();

        let counter = ctx.alloc_register();
        ctx.define(
            counter_name,
            VariableData {
                location: counter,
                is_const: false,
            },
        )?;

        let initial = compile_expr(ctx, initial)?;
        ctx.emit(inst![mov => counter, initial]);

        let end = try_err!(item.end(), "missing optional");
        let end = compile_expr(ctx, end)?;

        let step = if let Some(step) = item.step() {
            compile_expr(ctx, step)?
        } else {
            let step = ctx.alloc_register();
            let idx = ctx.alloc_constant(Value::from_int(1));
            ctx.emit(inst![cload => step, idx]);
            step
        };

        let label_body = ctx.new_label();
        ctx.set_label(label_body, ctx.next_offset());

        let block = try_err!(item.block(), "missing optional");
        compile_do(ctx, block)?;

        let tmp = ctx.alloc_register();
        ctx.emit(inst![add => counter, counter, step]);
        ctx.emit(inst![eq => tmp, counter, end]);
        ctx.emit(inst![isf => tmp]);
        ctx.emit(inst![cjmp => label_body]);
        Ok(())
    })
}

fn compile_for_gen(ctx: &mut Ctx, _item: ForGen) -> Result<(), Error> {
    scope_lexical_break(ctx, |_ctx| {
        panic!();
    })
}

#[cfg(test)]
mod tests {
    use std::{collections::BTreeMap, fs, ptr};

    use cstree::NodeCache;
    use insta::assert_snapshot;
    use paste::paste;

    use super::{super::disasm, compile, Heap, ObjectPtr};
    use crate::parser::{parse, syntax::Root};

    macro_rules! test {
        ($name:ident, $path:literal) => {
            paste! {
                #[test]
                fn [<test_compile_ $name>]() {
                    let mut cache = NodeCache::new();
                    let source = fs::read_to_string($path).unwrap();
                    let (syntax_tree, reports) = parse(&mut cache, &source);
                    assert!(reports.is_empty());
                    let root = Root::new(syntax_tree).unwrap();
                    let env = ObjectPtr::new(ptr::null_mut());
                    let heap = Heap::new();
                    let mut chunks = BTreeMap::new();
                    let chunk_alloc = &mut |chunk| {
                        let id = chunks.len() as u32;
                        chunks.insert(id, chunk);
                        id
                    };
                    let string_alloc: &mut dyn FnMut(&[u8]) -> ObjectPtr = &mut |bytes| heap.allocate_string(bytes);
                    compile(root, cache.interner(), 0, env, &heap, chunk_alloc, string_alloc).unwrap();
                    let mut asm = Vec::<u8>::new();
                    disasm::disasm(&mut asm, &mut chunks.values());
                    let asm_str = String::from_utf8(asm).unwrap();
                    assert_snapshot!(asm_str);
                }
            }
        };
    }

    test!(function, "../../test-files/function.lua");
    test!(if, "../../test-files/if.lua");
    test!(op_prec, "../../test-files/op_prec.lua");
    test!(nbody, "../../test-files/nbody.lua");
}