// COMPILER

use crate::parser::{
    Args, Block, Call, Chunk, Constructor, Expression, Function, LoadError, Name, Piece, Prefix,
    Statement, VarArg, Variable,
};
use crate::vm::{Code, Index, Machine, Offset, Prototype, Reference, Value};
use deimos::map::Map;
use std::cmp::PartialEq;
use std::ops::Deref;
use std::rc::Rc;

pub(crate) struct Traversal<'a, S> {
    // A 'cursor' of where we are in assumed execution
    machine: &'a mut Machine<S>,
    contexts: Vec<Context>, // Stack of function prototypes being compiled
}

pub(crate) struct Context {
    // Everything known about a closure at a time
    scopes: Vec<Scope>,
    vararg: Option<Piece<Option<String>>>,
    constants: Map<Value, Index>, // map some constant value to an index for not repeating constants
    prototype: Prototype,
}

#[derive(Clone, Debug)]
enum BlindJump {
    Break(usize),
    Continue(usize),
    Goto(String, usize),
}

#[derive(Debug)]
pub(crate) struct Scope {
    index: Index,          // Current index on stack
    jumps: Vec<BlindJump>, // todo: probably shouldn't have two vectors on each scope, they are used for temporaries too
    labels: Vec<(String, usize)>,
    symbol: Option<String>, // If that index has a symbol
    upvalue: bool,
}

impl<S> Traversal<'_, S> {
    fn compile<T>(&mut self, it: &T) -> Result<(), LoadError>
    where
        T: Compile,
    {
        it.compile(self)
    }
    fn push_code(&mut self, it: Piece<Code>) -> usize {
        match it.it {
            Code::TableNew => {
                self.temp_scope();
            }
            Code::TableSet => {
                self.consume_scope(3) // Value, Index, Table
            }
            Code::TableGet | Code::Binary(_) => {
                self.consume_scope(2); // Index, Table
                self.temp_scope(); // Result
            }
            Code::UpvalueGet(_) | Code::RegisterGet(_) | Code::Constant(_) | Code::Closure(_) => {
                self.temp_scope();
            }
            Code::Unary(_) => {
                self.consume_scope(1);
                self.temp_scope();
            }
            Code::TableInsert(_) => {
                self.consume_scope(2); // Value, Table
            }
            Code::Return(index) | Code::Call(index) => {
                let diff = self.index().unwrap() - index;
                self.consume_scope(diff + 1);
                //assert_eq!(self.last_scope().index, index);
                /* So note that Code::Call here doesn't do any temp_scope() magic, this is an edge-case for
                Code::Discard (if it's not used then Code::Call(index) will eat everything up to index
                and this consumption loop here won't do anything anyway */
            }
            Code::Discard(_index) => {
                // Unfortunately can't do any checks on the temporary scopes here, we don't know anything
                /* while self.last_scope().index > index { // Shouldn't this not even be here?
                  self.consume_scope(1);
                }
                while self.last_scope().index < index {
                  self.temp_scope();
                } */
            }
            Code::VarArg(count) => {
                for _ in 0..count {
                    self.temp_scope();
                }
                /* This is also weird like Code::Call and Code::Discard in that for count == 0 any number
                of vararg values are being pushed to the stack, but this is an edge-case that the compiler
                promises it will only use in Code::Return and Code::Call where they won't loop anyway */
            }
            Code::Skip | Code::UpvalueSet(_) | Code::RegisterSet(_) => {
                self.consume_scope(1);
            }
            Code::Close(_)
            | Code::Jump(_)
            | Code::RangeLoop(_)
            | Code::RangeBegin
            | Code::IteratorLoop(_)
            | Code::Message(_) => {}
        }
        let prototype = &mut self.contexts.last_mut().unwrap().prototype;
        prototype.code.push(it.it);
        prototype.pieces.push(it.unit());
        assert_eq!(prototype.code.len(), prototype.pieces.len());
        prototype.code.len() - 1
    }
    fn insert_code(&mut self, location: usize, it: Code) {
        // slightly dangerous since it doesn't consider scopes
        self.last_context().prototype.code[location] = it;
    }
    fn push_constant(&mut self, it: Value) -> Index {
        // todo: constants need pieces too!!!
        // don't use constants that already exist
        if let Some(entry) = self.contexts.last_mut().unwrap().constants.get(it.clone()) {
            *entry
        } else {
            let index = {
                let constants = &mut self.contexts.last_mut().unwrap().prototype.constants;
                constants.push(it.clone());
                constants.len() - 1
            }
            .try_into()
            .unwrap();
            self.contexts.last_mut().unwrap().constants.set(it, index);
            index
        }
    }
    fn push_literal(&mut self, it: Value, piece: Piece<()>) {
        let constant = self.push_constant(it);
        self.push_code(piece.swap(Code::Constant(constant)));
    }
    fn push_string(&mut self, string: String, piece: Piece<()>) -> Result<(), LoadError> {
        match self.machine.new_string(string.clone().as_str()) {
            Err(error) => Err(LoadError(error.0, piece)),
            Ok(string) => Ok(self.push_literal(string, piece)),
        }
    }
    fn new_scope(&mut self) {
        // Make some new scope
        let index = self.last_scope().index;
        self.contexts.last_mut().unwrap().scopes.push(Scope {
            index,
            jumps: vec![],
            labels: vec![],
            symbol: None,
            upvalue: false,
        });
    }
    fn temp_scope(&mut self) -> Index {
        // Make a new scope with index++
        self.new_scope();
        if self.last_context().scopes.len() > 2 {
            // For the dummy scope
            self.last_scope().index += 1;
        }
        self.last_scope().index
    }
    fn last_scope(&mut self) -> &mut Scope {
        self.last_context().scopes.last_mut().unwrap()
    }
    fn index(&mut self) -> Option<Index> {
        if self.last_context().scopes.len() == 1 {
            None
        } else {
            Some(self.last_context().scopes.last_mut().unwrap().index)
        }
    }
    fn code_length(&mut self) -> usize {
        self.last_context().prototype.code.len()
    }
    fn close_scope(&mut self) {
        // Handle exiting a variable register's scope
        let mut last = self.last_context().scopes.pop().unwrap(); // Remove it
        self.last_scope().jumps.append(&mut last.jumps); // Bubble up the orphaned blind jumps
        // If this scope's declaration was an upvalue then the VM needs to close it at runtime
        assert!(last.symbol.is_some());
        // Always have to pop off variable registers
        if let Some(Code::Close(_)) = self.last_context().prototype.code.last() {
            // Did we just close a scope?
            let len = self.code_length(); // Ok, what's the last instruction index?
            self.insert_code(len - 1, Code::Close(last.index)); // Push that Code::Close even further back to this register.
        }
    }
    fn consume_scope(&mut self, times: u16) {
        // Handle consuming a temporary on the stack
        for _ in 0..times {
            assert!(self.last_context().scopes.pop().unwrap().symbol.is_none()); // Remove it
        }
    }
    fn last_context(&mut self) -> &mut Context {
        self.contexts.last_mut().unwrap()
    }
    pub(crate) fn close_context(&mut self) -> Result<Prototype, LoadError> {
        let mut context = self.contexts.pop().unwrap();
        for jump in context.scopes.pop().unwrap().jumps {
            // Get back the dummy one and error any jumps
            return Err(LoadError::from_str(
                format!("{:?}", jump).as_str(),
                &Piece::<()>::null(),
            )); // todo: we need Vec or some append in LoadError and also BlindJumps need pieces!
        }
        Ok(context.prototype)
    }
    pub(crate) fn new(machine: &mut Machine<S>) -> Traversal<S> {
        Traversal {
            machine,
            contexts: vec![Context {
                scopes: vec![Scope {
                    // Dummy scope
                    index: 0,
                    jumps: vec![],
                    labels: vec![],
                    symbol: None,
                    upvalue: false,
                }],
                vararg: Some(Piece::<Name>::null().swap(None)),
                constants: Map::new(),
                prototype: Prototype {
                    args: 0,
                    vararg: true,
                    prototypes: vec![],
                    constants: vec![],
                    upvalues: vec![],
                    pieces: vec![],
                    code: vec![],
                },
            }],
        }
    }
}

// todo: not all impl Compile for Piece<T> needs to be &T
trait Compile {
    fn compile<S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError>;
}

fn compile_function<S>(
    this: &Piece<&Function>,
    traversal: &mut Traversal<S>,
    method: bool,
) -> Result<(), LoadError> {
    let mut self_arguments = Vec::new();
    if method {
        self_arguments.push("self".to_string());
    }
    self_arguments.append(&mut this.it.args.iter().map(|it| it.it.0.clone()).collect());
    let mut argument_scopes = vec![Scope {
        // Dummy initial scope so the rest can start from index = 0 (since there can be zero arguments!!)
        index: 0,
        jumps: vec![],
        labels: vec![],
        symbol: None,
        upvalue: false,
    }];
    let mut index = 0;
    for argument in self_arguments.iter().rev() {
        argument_scopes.push(Scope {
            index,
            jumps: vec![],
            labels: vec![],
            symbol: Some(argument.clone()),
            upvalue: false,
        });
        index += 1;
    }
    traversal.contexts.push(Context {
        vararg: this
            .it
            .vararg
            .as_ref()
            .map(|it| it.swap(it.clone().it.0.map(|string| string.0.clone()))),
        scopes: argument_scopes,
        prototype: Prototype {
            args: this.it.args.len().try_into().unwrap(),
            vararg: this.it.vararg.is_some(),
            prototypes: vec![],
            constants: vec![],
            upvalues: vec![Reference::Upvalue(0)], // _ENV
            pieces: vec![],
            code: vec![],
        },
        constants: Map::new(),
    });
    traversal.compile(&this.it.block.inner_ref())?; // compile the block
    let prototype = traversal.close_context()?;
    let prototypes = &mut traversal.last_context().prototype.prototypes; // add this to the upper function's prototypes list
    let index = prototypes.len().try_into().unwrap();
    prototypes.push(Rc::new(prototype));
    traversal.push_code(this.swap(Code::Closure(index))); // reference this closure
    Ok(())
}

enum Access {
    Set,
    Get,
}

fn compile_identifier<S>(
    name: Piece<String>,
    traversal: &mut Traversal<S>,
    access: Access,
) -> Result<(), LoadError> {
    let contexts_len = traversal.contexts.len();
    for context_index in 0..contexts_len {
        // For each context
        let context_index = contexts_len - context_index - 1; // From top to bottom
        if context_index == contexts_len - 1 {
            // If this is our context
            for scope in traversal.contexts[context_index].scopes.iter().rev() {
                // Go through the scopes
                if scope.symbol.as_ref().is_some_and(|it| it.eq(&name.it)) {
                    // Find it
                    let index = scope.index;
                    traversal.push_code(name.swap(match access {
                        Access::Set => Code::RegisterSet(index),
                        Access::Get => Code::RegisterGet(index),
                    }));
                    return Ok(()); // Then break out
                }
            }
        } else {
            // Otherwise it's another function's 'scope'(s)
            let scopes_len = traversal.contexts[context_index].scopes.len();
            for scope_index in 0..scopes_len {
                // For each scope in that foreign context
                let scope_index = scopes_len - scope_index - 1; // From top to bottom
                if traversal.contexts[context_index].scopes[scope_index]
                    .symbol
                    .as_ref()
                    .is_some_and(|it| it.eq(&name.it))
                {
                    let register = traversal.contexts[context_index].scopes[scope_index].index;
                    traversal.contexts[context_index].scopes[scope_index].upvalue = true;
                    let mut reference = Reference::Closing(register);
                    // its important that a few conditions are held here e.g. length > 2, j < length + 1...
                    for down_context_index in context_index + 1..contexts_len {
                        // Ok now going from every context to the one wanting the upvalue
                        // set reference to what we found for the next context
                        reference = match traversal.contexts[down_context_index]
                            .prototype
                            .upvalues
                            .iter()
                            .position(|it| *it == reference)
                        {
                            Some(index) => Reference::Upvalue(index.try_into().unwrap()),
                            None => {
                                // we don't even have a reference to this in our upvalues so add it
                                let length = traversal.contexts[down_context_index]
                                    .prototype
                                    .upvalues
                                    .len();
                                traversal.contexts[down_context_index]
                                    .prototype
                                    .upvalues
                                    .push(reference);
                                Reference::Upvalue(length.try_into().unwrap()) // Reference (child) to this reference to that reference (parent) for next context
                            }
                        }
                    }
                    let index = match reference {
                        Reference::Closing(_) => panic!(),
                        Reference::Upvalue(index) => index,
                    };
                    traversal.push_code(name.swap(match access {
                        Access::Set => Code::UpvalueSet(index),
                        Access::Get => Code::UpvalueGet(index),
                    }));
                    return Ok(());
                }
            }
        }
    }
    // At this point there were no returns so we must use _G table
    traversal.push_string(name.it.clone(), name.unit())?;
    traversal.push_code(name.swap(Code::UpvalueGet(0)));
    traversal.push_code(name.swap(match access {
        Access::Set => Code::TableSet,
        Access::Get => Code::TableGet,
    }));
    Ok(())
}

fn compile_variable<S>(
    variable: Piece<&Variable>,
    traversal: &mut Traversal<S>,
    access: Access,
) -> Result<(), LoadError> {
    match &variable.it {
        Variable::Name(Name(name)) => {
            compile_identifier(variable.swap(name.clone()), traversal, access)?;
        }
        Variable::Index(prefix, index) => {
            // Compile the prefix expression, the index, and set
            traversal.compile(&index.inner_ref())?;
            traversal.compile(&prefix.as_ref().inner_ref())?;
            traversal.push_code(prefix.unit().end().swap(match access {
                Access::Set => Code::TableSet,
                Access::Get => Code::TableGet,
            }));
        }
        Variable::Register(index) => {
            traversal.push_code(variable.swap(match access {
                Access::Set => Code::RegisterSet(*index),
                Access::Get => Code::RegisterGet(*index),
            }));
        }
    }
    Ok(())
}

impl Compile for Piece<&Function> {
    fn compile<S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        compile_function(self, traversal, false)
    }
}

impl Compile for Piece<&Args> {
    fn compile<S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        let mut expressions_iter = self.it.0.iter().rev();
        if let Some(expression) = expressions_iter.next() {
            let piece = expression.unit();
            match &expression.it {
                Expression::Prefix(prefix) => match &**prefix {
                    Prefix::Call(call) => {
                        traversal.compile(&piece.swap(call))?;
                    }
                    prefix => {
                        traversal.compile(&piece.swap(prefix))?;
                    }
                },
                Expression::VarArg(vararg) => {
                    traversal.compile(&piece.swap(vararg))?;
                    traversal.push_code(piece.swap(Code::VarArg(0))); // Dump all the varargs
                }
                _ => {
                    traversal.compile(&expression.inner_ref())?;
                }
            }
        }
        while let Some(expression) = expressions_iter.next() {
            traversal.compile(&expression.inner_ref())?;
        }
        Ok(())
    }
}

impl Compile for Piece<&Call> {
    fn compile<S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        let index = traversal.index().map(|it| it + 1).unwrap_or(0); // there will always be at least the callable
        traversal.compile(&self.it.args.inner_ref())?; // arguments
        traversal.compile(&self.it.prefix.as_ref().inner_ref())?; // call prefix
        match &self.it.method {
            Some(name) => {
                traversal.push_string(name.it.0.clone(), name.unit())?;
                traversal.push_code(name.swap(Code::TableGet));
            }
            None => {}
        }
        traversal.push_code(self.it.prefix.unit().end().swap(Code::Call(index))); // Our responsibility to discard here
        Ok(())
    }
}

impl Compile for Piece<&Prefix> {
    fn compile<'a, 'b, S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        match &self.it {
            Prefix::Call(call) => {
                let index = traversal.index().unwrap();
                traversal.compile(&self.swap(call))?;
                traversal.push_code(self.swap(Code::Discard(index + 1)).end()); // produce 1 return value, or pad nil
            }
            Prefix::Variable(variable) => {
                compile_variable(self.swap(variable), traversal, Access::Get)?;
            }
            Prefix::Expression(expression) => {
                traversal.compile(&self.swap(expression))?;
            }
        }
        Ok(())
    }
}

// for the sake of erroring alone
impl Compile for Piece<&VarArg> {
    fn compile<'a, 'b, S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        match &traversal.contexts.last().unwrap().vararg {
            Some(piece) => match &piece.it {
                Some(that_name) => match &self.it.0 {
                    Some(Name(this_name)) => {
                        if that_name.eq(this_name) {
                            Ok(())
                        } else {
                            Err(LoadError(
                                format!(
                                    "This context's only visible vararg parameter '{:?}' is not named '{:?}'",
                                    that_name, this_name
                                ),
                                self.unit(),
                            ))
                        }
                    }
                    None => Err(LoadError(
                        format!(
                            "This context's only visible vararg parameter must be referred to by its name '{:?}'",
                            that_name
                        ),
                        self.unit(),
                    )),
                },
                None => match &self.it.0 {
                    Some(name) => Err(LoadError(
                        format!(
                            "This context's only visible vararg parameter has no name '{:?}'",
                            name
                        ),
                        self.unit(),
                    )),
                    None => Ok(()),
                },
            },
            None => Err(LoadError(
                format!(
                    "This context has no accessible vararg parameter '{:?}'",
                    self.it
                ),
                self.unit(),
            )),
        }
    }
}

impl Compile for Piece<&Constructor> {
    fn compile<'a, 'b, S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        traversal.push_code(self.swap(Code::TableNew));
        let index_table = traversal.index().unwrap();
        for field in self.it.0.iter() {
            match &field.it.index {
                Some(index) => {
                    traversal.compile(&index.inner_ref())?;
                    traversal.push_code(
                        field
                            .it
                            .expression
                            .swap(Code::RegisterGet(index_table))
                            .start(),
                    );
                    traversal.push_code(field.it.expression.swap(Code::TableSet));
                }
                None => {
                    match &field.it.expression.it {
                        Expression::VarArg(vararg) => {
                            traversal.compile(&field.it.expression.swap(vararg))?;
                            traversal.push_code(field.it.expression.swap(Code::VarArg(0)));
                        }
                        it => {
                            traversal.compile(&field.it.expression.swap(it))?;
                        }
                    }
                    let index = traversal
                        .contexts
                        .last()
                        .unwrap()
                        .scopes
                        .last()
                        .unwrap()
                        .index;
                    traversal.push_code(
                        field
                            .it
                            .expression
                            .swap(Code::RegisterGet(index_table))
                            .start(),
                    );
                    traversal.push_code(field.it.expression.swap(Code::TableInsert(index)));
                }
            }
        }
        Ok(())
    }
}

impl Compile for Piece<&Expression> {
    fn compile<'a, S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        match self.it {
            Expression::Nil => {
                traversal.push_literal(Value::Nil, self.unit());
            }
            Expression::True => {
                traversal.push_literal(Value::Bool(true), self.unit());
            }
            Expression::False => {
                traversal.push_literal(Value::Bool(false), self.unit());
            }
            Expression::Number(number) => {
                let number = number
                    .0
                    .parse::<u64>()
                    .map(|it| Value::Integer(it))
                    .unwrap_or(Value::Number(
                        number
                            .0
                            .parse::<f64>()
                            .map_err(|it| LoadError(it.to_string(), self.unit()))?,
                    ));
                traversal.push_literal(number, self.unit());
            }
            Expression::String(string) => {
                traversal.push_string(string.clone(), self.unit())?;
            }
            Expression::VarArg(vararg) => {
                traversal.compile(&self.swap(vararg))?;
                traversal.push_code(self.swap(Code::VarArg(1)));
            }
            Expression::Function(function) => {
                traversal.compile(&self.swap(function))?;
            }
            Expression::Prefix(prefix) => {
                traversal.compile(&self.swap(prefix.as_ref()))?;
            }
            Expression::Table(constructor) => {
                traversal.compile(&self.swap(constructor))?;
            }
            Expression::Binary(left, op, right) => {
                traversal.compile(&left.as_ref().inner_ref())?;
                traversal.compile(&right.as_ref().inner_ref())?;
                traversal.push_code(op.swap(Code::Binary(op.it)));
            }
            Expression::Unary(op, expression) => {
                traversal.compile(&expression.as_ref().inner_ref())?;
                traversal.push_code(op.swap(Code::Unary(op.it)));
            }
        };
        Ok(())
    }
}

fn jump_from(here: usize, there: usize) -> Result<Offset, LoadError> {
    Ok((there as i64 - here as i64) as Offset)
}

fn resolve_break<S>(
    traversal: &mut Traversal<S>,
    begin: usize,
    end: usize,
) -> Result<(), LoadError> {
    for jump in std::mem::replace(&mut traversal.last_scope().jumps, Vec::new()) {
        // ?
        match jump {
            BlindJump::Continue(index) => {
                traversal.insert_code(index, Code::Jump(jump_from(index, begin)?));
            }
            BlindJump::Break(index) => {
                traversal.insert_code(index, Code::Jump(jump_from(index, end)?));
            }
            _ => traversal.last_scope().jumps.push(jump),
        }
    }
    Ok(())
}

fn compile_assignment<S>(
    traversal: &mut Traversal<S>,
    variables: &Vec<Piece<Variable>>,
    expressions: &Vec<Piece<Expression>>,
) -> Result<(), LoadError> {
    let mut variables_iter = variables.iter().peekable();
    let mut expressions_iter = expressions.iter().peekable();
    while let Some(expression) = expressions_iter.next() {
        let compile_variable_expression = |// Compile a variable normally with one value from an expression
                                           traversal: &mut Traversal<S>,
                                           variable: Option<&Piece<Variable>>,
                                           expression: &Piece<&Expression>|
         -> Result<(), LoadError> {
            let index = traversal.index().unwrap_or(0);
            traversal.compile(expression)?;
            if let Some(variable) = variable {
                compile_variable(variable.inner_ref(), traversal, Access::Set)?;
            } else {
                traversal.push_code(expression.swap(Code::Discard(index)));
            }
            Ok(())
        };
        if expressions_iter.peek().is_none() {
            // If this is the last expression it CAN have multiplicity
            match &expression.it {
                Expression::Prefix(prefix) => {
                    match **prefix {
                        Prefix::Call(ref call) => {
                            // multiple
                            let mut index = traversal.index().unwrap();
                            traversal.compile(&expression.swap(call))?;
                            let index_discard = traversal.push_code(expression.swap(Code::UNSURE));
                            while let Some(variable) = variables_iter.next() {
                                index += 1; // The index we would need to pad nil to or restrict to
                                compile_variable(variable.inner_ref(), traversal, Access::Set)?; // Pops each return value off the stack
                            }
                            traversal.insert_code(
                                index_discard,
                                Code::Discard(index.try_into().unwrap()),
                            );
                        }
                        ref prefix => {
                            if let Some(variable) = variables_iter.next() {
                                traversal.compile(&expression.swap(prefix))?;
                                compile_variable(variable.inner_ref(), traversal, Access::Set)?
                            }
                        }
                    }
                }
                Expression::VarArg(vararg) => {
                    traversal.compile(&expression.swap(vararg))?; // Does nothing, error check
                    if variables_iter.peek().is_some() {
                        let mut count = 0;
                        let index_vararg = traversal.push_code(expression.swap(Code::VarArg(0)));
                        for variable in variables_iter.clone() {
                            count += 1;
                            compile_variable(variable.inner_ref(), traversal, Access::Set)?;
                        }
                        traversal
                            .insert_code(index_vararg, Code::VarArg(count.try_into().unwrap()));
                    }
                }
                _ => {
                    compile_variable_expression(
                        traversal,
                        variables_iter.next(),
                        &expression.inner_ref(),
                    )?;
                }
            }
        } else {
            // otherwise treat it normally and consume one variable
            compile_variable_expression(traversal, variables_iter.next(), &expression.inner_ref())?;
        }
    }
    while let Some(variable) = variables_iter.next() {
        // Set the rest of the variables to nil
        traversal.push_literal(Value::Nil, variable.unit());
        compile_variable(variable.inner_ref(), traversal, Access::Set)?;
    }
    Ok(())
}

impl Compile for Piece<&Block> {
    fn compile<S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        let last_scope_index = traversal.index().unwrap_or(0);
        let mut depth = 0; // Counter for 'dangling' scopes produced by any declarations to be closed.
        // todo: for some reason, register indices are called "<thing>_index" and code indices "index_<thing>"? fix it
        // todo: maybe refactor this so that there are less indents
        for statement in self.it.statements.iter() {
            let piece = statement.unit();
            match &statement.it {
                Statement::Semicolon => {}
                Statement::Assignment(variables, expressions) => {
                    compile_assignment(traversal, variables, expressions)?;
                }
                Statement::Call(call) => {
                    let index = traversal.index().unwrap_or(0);
                    traversal.compile(&self.swap(call))?;
                    traversal.push_code(self.swap(Code::Discard(index)));
                }
                Statement::Label(label) => {
                    let index_label = traversal.code_length() - 1;
                    traversal
                        .last_scope()
                        .labels
                        .push((label.it.0.clone(), index_label));
                    let mut jumps = Vec::new();
                    for jump in traversal.last_scope().jumps.clone().iter().rev() {
                        match jump {
                            BlindJump::Goto(name, index) => {
                                if name.eq(&label.it.0) {
                                    traversal.insert_code(
                                        *index,
                                        Code::Jump(jump_from(*index, index_label)?),
                                    )
                                } else {
                                    jumps.push(jump);
                                }
                            }
                            _ => jumps.push(jump),
                        }
                    }
                }
                Statement::Break => {
                    let index = traversal.push_code(piece.swap(Code::UNSURE));
                    traversal.last_scope().jumps.push(BlindJump::Break(index));
                }
                Statement::Continue => {
                    let index = traversal.push_code(piece.swap(Code::UNSURE));
                    traversal
                        .last_scope()
                        .jumps
                        .push(BlindJump::Continue(index));
                }
                Statement::Goto(name) => {
                    let mut index = None;
                    for scope in traversal.contexts.last().unwrap().scopes.iter().rev() {
                        for (label_name, label_index) in scope.labels.iter().rev() {
                            if name.it.0 == *label_name {
                                index = Some(label_index.clone());
                                break;
                            }
                        }
                    }
                    if let Some(index) = index {
                        let current = (traversal.last_context().prototype.code.len() - 1) as Offset;
                        traversal.push_code(piece.swap(Code::Jump(index as Offset - current)));
                    } else {
                        let index = traversal.push_code(piece.swap(Code::UNSURE));
                        traversal
                            .last_scope()
                            .jumps
                            .push(BlindJump::Goto(name.it.0.clone(), index));
                    }
                }
                Statement::Do(block) => {
                    traversal.compile(&self.swap(block))?;
                }
                Statement::While(condition, block) => {
                    let piece = condition.unit().end();
                    let index_exp = traversal.code_length();
                    traversal.compile(&condition.inner_ref())?;
                    traversal.push_code(piece.swap(Code::Skip));
                    let index_exit = traversal.push_code(piece.swap(Code::UNSURE));
                    traversal.compile(&self.swap(block))?;
                    let index_jump = traversal.code_length();
                    resolve_break(traversal, index_exp, index_jump + 1)?;
                    traversal.push_code(piece.swap(Code::Jump(jump_from(index_jump, index_exp)?)));
                    traversal.insert_code(
                        index_exit,
                        Code::Jump(jump_from(index_exit, index_jump + 1)?),
                    );
                }
                Statement::Repeat(block, condition) => {
                    let piece = condition.unit().end();
                    let index_block = traversal.code_length();
                    traversal.compile(&self.swap(block))?;
                    let index_exp = traversal.code_length();
                    traversal.compile(&condition.inner_ref())?;
                    traversal.push_code(piece.swap(Code::Skip));
                    let index_jump = traversal.code_length();
                    resolve_break(traversal, index_exp, index_jump + 1)?;
                    traversal
                        .push_code(piece.swap(Code::Jump(jump_from(index_jump, index_block)?)));
                }
                Statement::If(conditions, blocks, otherwise) => {
                    let mut conditions = conditions.iter().peekable();
                    let mut blocks = blocks.iter();
                    assert_eq!(conditions.len(), blocks.len());
                    let mut jumps = Vec::new();
                    while let Some(condition) = conditions.next() {
                        let block = blocks.next().unwrap();
                        traversal.compile(&condition.inner_ref())?;
                        traversal.push_code(condition.swap(Code::Skip));
                        let index_next = traversal.push_code(condition.swap(Code::UNSURE)); // jump to the next branch
                        traversal.compile(&block.inner_ref())?;
                        if conditions.peek().is_some() || otherwise.is_some() {
                            // don't jump ahead by 1 if it's at the end already and there is no else block
                            jumps.push(traversal.push_code(condition.swap(Code::UNSURE))); // jump all the way to the end
                        }
                        let length = traversal.code_length();
                        traversal
                            .insert_code(index_next, Code::Jump(jump_from(index_next, length)?));
                    }
                    assert!(blocks.next().is_none());
                    if let Some(otherwise) = otherwise {
                        traversal.compile(&otherwise.inner_ref())?;
                    }
                    for jump in jumps {
                        let length = traversal.code_length();
                        traversal.insert_code(jump, Code::Jump(jump_from(jump, length)?))
                    }
                }
                Statement::Range(var, init, limit, step, block) => {
                    traversal.compile(&init.inner_ref())?;
                    traversal.compile(&limit.inner_ref())?;
                    match step {
                        Some(expression) => traversal.compile(&expression.inner_ref())?,
                        None => traversal.push_literal(Value::Integer(1), limit.unit().end()),
                    };
                    let index_loop =
                        traversal.push_code(block.unit().start().swap(Code::RangeBegin)) + 1;
                    traversal.temp_scope();
                    traversal.last_scope().symbol = Some(var.it.0.clone());
                    traversal.compile(&block.inner_ref())?;
                    traversal.close_scope(); // temporary for v
                    let index_jump = traversal.push_code(block.unit().start().swap(Code::UNSURE)); // Dummy NOP
                    traversal.insert_code(
                        index_jump,
                        Code::RangeLoop(jump_from(index_jump, index_loop)?),
                    ); // Reinsert into above
                    resolve_break(traversal, index_loop, index_jump + 1)?;
                    traversal.consume_scope(3); // init, limit, step
                }
                Statement::Iterator(names, expressions, block) => {
                    let mut iterator_variables = Vec::new();
                    for _ in 0..2 {
                        // f, s, var
                        traversal.new_scope();
                        iterator_variables.push(
                            block
                                .swap(Variable::Register(traversal.index().unwrap()))
                                .start(),
                        );
                    }
                    // So we put 'explist' into f, s, var as in https://www.lua.org/manual/5.3/manual.html#3.3.5
                    compile_assignment(traversal, &iterator_variables, expressions)?;
                    let index_loop = traversal.push_code(block.swap(Code::UNSURE));
                    // Discard to fill all names
                    let index = traversal.index().unwrap();
                    traversal.push_code(
                        block
                            .swap(Code::Discard(
                                (index as usize + names.len()).try_into().unwrap(),
                            ))
                            .start(),
                    );
                    for name in names.iter() {
                        traversal.new_scope();
                        traversal.last_scope().symbol = Some(name.it.0.clone());
                    }
                    traversal.compile(&block.inner_ref())?;
                    for _name in names.iter() {
                        traversal.close_scope();
                    }
                    let index_jump = traversal.push_code(block.unit().end().swap(Code::UNSURE));
                    traversal
                        .insert_code(index_jump, Code::Jump(jump_from(index_jump, index_loop)?));
                    traversal.insert_code(
                        index_loop,
                        Code::IteratorLoop(jump_from(index_loop, index_jump + 1)?),
                    );
                    resolve_break(traversal, index_loop, index_jump + 1)?;
                    traversal.close_scope();
                    traversal.close_scope();
                    traversal.close_scope(); // f, s, var
                }
                Statement::Function(names, method_name, function) => {
                    // todo: ideally make names an expression, change the parser for that too
                    let mut names_method = names.clone();
                    let is_method = if let Some(method) = method_name {
                        // for use in compiling the function with 'self' as an argument
                        names_method.push(method.clone());
                        true
                    } else {
                        false
                    };
                    let mut names = names_method.iter().peekable();
                    let first = names.next().unwrap();
                    let var = first.swap(Variable::Name(first.it.clone()));
                    if !names.peek().is_some() {
                        compile_function(&self.swap(function), traversal, is_method)?;
                        compile_variable(var.inner_ref(), traversal, Access::Set)?;
                    } else {
                        compile_variable(var.inner_ref(), traversal, Access::Get)?;
                        while let Some(name) = names.next() {
                            if !names.peek().is_some() {
                                compile_function(&self.swap(function), traversal, is_method)?;
                                traversal.push_string(name.it.0.clone(), name.unit())?;
                                traversal.push_code(name.unit().end().swap(Code::TableSet));
                            } else {
                                traversal.push_string(name.it.0.clone(), name.unit())?;
                                traversal.push_code(name.unit().end().swap(Code::TableGet));
                            }
                        }
                    }
                }
                Statement::LocalFunction(name, function) => {
                    traversal.compile(&self.swap(function))?;
                    traversal.last_scope().symbol = Some(name.it.0.clone());
                    depth += 1;
                }
                Statement::Declaration(names, expressions) => {
                    let mut expressions_iter = expressions.iter().peekable();
                    let mut names_iter = names.iter().peekable();
                    let compile_declaration = |traversal: &mut Traversal<S>,
                                               expression: &Piece<Expression>,
                                               name: &Piece<Name>|
                     -> Result<(), LoadError> {
                        traversal.compile(&expression.inner_ref())?;
                        traversal.last_scope().symbol = Some(name.it.0.clone());
                        Ok(())
                    };
                    while names_iter.peek().is_some() {
                        match (expressions_iter.next(), expressions_iter.peek().is_none()) {
                            (Some(expression), true) => {
                                // expression is last, has multiplicity
                                match expression.it {
                                    Expression::VarArg(ref vararg) => {
                                        let names: Vec<&Piece<Name>> = names_iter.collect();
                                        traversal.compile(&expression.swap(vararg))?; // for error checking only
                                        traversal.push_code(
                                            expression.swap(Code::VarArg(
                                                names.len().try_into().unwrap(),
                                            )),
                                        );
                                        let mut scopes =
                                            traversal.last_context().scopes.iter_mut().rev();
                                        for name in names.iter().rev() {
                                            scopes.next().unwrap().symbol = Some(name.it.0.clone());
                                            depth += 1;
                                        }
                                    }
                                    Expression::Prefix(ref prefix) => {
                                        match prefix.deref() {
                                            Prefix::Call(call) => {
                                                let names: Vec<&Piece<Name>> = names_iter.collect();
                                                traversal.compile(&expression.swap(call))?;
                                                let index = (traversal.index().unwrap_or(0)
                                                    as usize
                                                    + names.len())
                                                .try_into()
                                                .unwrap();
                                                traversal.push_code(
                                                    expression.swap(Code::Discard(index)).end(),
                                                ); // discard and extend to names
                                                let mut scopes = traversal
                                                    .last_context()
                                                    .scopes
                                                    .iter_mut()
                                                    .rev();
                                                for name in names.iter().rev() {
                                                    scopes.next().unwrap().symbol =
                                                        Some(name.it.0.clone());
                                                    depth += 1;
                                                }
                                            }
                                            _ => {
                                                compile_declaration(
                                                    traversal,
                                                    expression,
                                                    names_iter.next().unwrap(),
                                                )?;
                                                depth += 1;
                                            }
                                        }
                                    }
                                    _ => {
                                        compile_declaration(
                                            traversal,
                                            expression,
                                            names_iter.next().unwrap(),
                                        )?;
                                        depth += 1;
                                    }
                                }
                                break;
                            }
                            (Some(expression), false) => {
                                compile_declaration(
                                    traversal,
                                    expression,
                                    names_iter.next().unwrap(),
                                )?;
                                depth += 1;
                            } // expression is not last, just use it
                            (None, ..) => {
                                // no more expressions
                                let index = traversal.index().unwrap_or(0) as usize;
                                traversal.push_code(
                                    names_iter
                                        .next()
                                        .unwrap()
                                        .swap(Code::Discard(
                                            (index + names_iter.count()).try_into().unwrap(),
                                        ))
                                        .end(),
                                );
                                break;
                            }
                        }
                    }
                    let index = traversal.index().unwrap_or(0);
                    // clear up the remaining expressions and discard if they existed
                    if expressions_iter.peek().is_some() {
                        while let Some(expression) = expressions_iter.next() {
                            traversal.compile(&expression.inner_ref())?
                        }
                        traversal.push_code(Piece::<()>::null().swap(Code::Discard(index)));
                    }
                }
            }
        }
        // Any 'dangling' scopes that weren't closed (declarations) will be at the end of this block
        let return_index = traversal.index().unwrap_or(0);
        if let Some(return_expressions) = &self.it.return_exp {
            // arguments are backwards
            for return_expression in return_expressions.iter().rev() {
                traversal.compile(&return_expression.inner_ref())?;
            }
            if !return_expressions.is_empty() {
                traversal.push_code(
                    return_expressions
                        .first()
                        .unwrap()
                        .swap(Code::Return(return_index)),
                );
            }
        }
        for _ in 0..depth {
            traversal.close_scope();
        }
        assert_eq!(traversal.index().unwrap_or(0), last_scope_index);
        Ok(())
    }
}

impl Compile for Chunk {
    fn compile<S>(&self, traversal: &mut Traversal<S>) -> Result<(), LoadError> {
        traversal.compile(&self.0.inner_ref())
    }
}

pub(crate) fn compile<S>(chunk: &Chunk, machine: &mut Machine<S>) -> Result<Prototype, LoadError> {
    let mut traversal = Traversal::new(machine);
    traversal.compile(chunk)?;
    traversal.close_context()
}