lysp/src/compile/block.rs

use std::{
    collections::HashMap,
    ops::{Deref, DerefMut},
    rc::Rc,
};

use crate::{
    compile::CompileOptions,
    vm::value::{BytecodeFunction, StringValue},
};
use crate::{
    compile::{
        error::CompileError,
        function::FunctionSignature,
        instruction::Emitted,
        syntax::{Expression, FunctionBody},
    },
    vm::{
        Value,
        instruction::{ConstantId, ImmediateInteger, Instruction, LocalId},
        value::{BooleanValue, IdentifierValue, NumberValue},
    },
};

#[derive(Clone)]
pub struct FunctionBlock {
    parent: Option<usize>,
    identifier: Option<IdentifierValue>,
    docstring: Option<StringValue>,
    signature: FunctionSignature,

    // Data
    pub(crate) constants: Vec<Value>,
    locals: Vec<Local>,
    upvalues: Vec<UpvalueDef>,
    scope_depth: usize,

    // Code
    pub(crate) instructions: Vec<Emitted>,
    labels: Vec<usize>,
    loop_stack: Vec<LoopStackEntry>,
}

#[derive(Clone)]
struct Local {
    name: IdentifierValue,
    depth: isize,
    is_captured: bool,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct UpvalueDef {
    pub(crate) index: LocalId,
    pub(crate) is_local: bool,
}

#[derive(Clone)]
struct LoopStackEntry {
    label_entry: usize,
    label_exit: usize,
}

#[derive(Debug, PartialEq)]
pub enum CompileValue {
    Nil,
    Identifier(IdentifierValue),
    Integer(NumberValue),
    Boolean(BooleanValue),
    String(StringValue),
    LocalFunction(usize),
    Quote(Rc<Value>),
    Value(Value),
    Stack,
}

#[derive(Clone)]
pub struct CompileContext {
    pub function_blocks: Vec<FunctionBlock>,
    pub(crate) current: usize,
    pub(crate) options: CompileOptions,
}

pub trait Compile {
    fn compile(&self, cx: &mut CompileContext) -> Result<CompileValue, CompileError>;
}

impl CompileContext {
    pub fn new(options: CompileOptions, root_name: Option<IdentifierValue>) -> Self {
        Self {
            function_blocks: vec![FunctionBlock::root(root_name)],
            current: 0,
            options,
        }
    }

    pub fn compile_value(
        options: CompileOptions,
        chunk_name: Option<IdentifierValue>,
        value: &Value,
    ) -> Result<Rc<BytecodeFunction>, CompileError> {
        let mut cx = Self::new(options, chunk_name);
        let expression = Expression::parse(value).map_err(CompileError::Parse)?;
        let value = expression.compile(&mut cx)?;
        cx.compile_return_value(value)?;
        cx.to_bytecode()
    }

    pub fn compile_function(
        &mut self,
        identifier: Option<IdentifierValue>,
        docstring: Option<StringValue>,
        signature: &FunctionSignature,
        body: &FunctionBody,
    ) -> Result<usize, CompileError> {
        // TODO signature
        let index = self.push_lambda_context(identifier, docstring, signature)?;
        for expression in body.head.iter() {
            self.compile_statement(expression)?;
        }
        let value = body.tail.compile(self)?;
        self.compile_return_value(value)?;
        self.pop_context();
        Ok(index)
    }

    pub fn compile_return_value(&mut self, value: CompileValue) -> Result<(), CompileError> {
        self.push(value)?;
        self.emit(Instruction::Return);
        Ok(())
    }

    pub fn push_lambda_context(
        &mut self,
        identifier: Option<IdentifierValue>,
        docstring: Option<StringValue>,
        signature: &FunctionSignature,
    ) -> Result<usize, CompileError> {
        if self.options.trace_compile {
            eprintln!("COMPILE: push_lambda_context({identifier:?})");
        }
        let block = FunctionBlock::new(Some(self.current), identifier, docstring, signature);
        let index = self.function_blocks.len();
        self.function_blocks.push(block);
        self.current = index;
        Ok(index)
    }

    pub fn pop_context(&mut self) {
        if self.options.trace_compile {
            let current_name = self.identifier.clone();
            eprintln!("COMPILE: pop_lambda_context({current_name:?})");
        }
        let index = self
            .parent
            .expect("cannot pop out of root function context");
        self.current = index;
    }

    pub fn compile_statement(&mut self, expression: &Expression) -> Result<(), CompileError> {
        let value = expression.compile(self)?;
        self.discard(value);
        Ok(())
    }

    pub fn compile_set_global(
        &mut self,
        identifier: IdentifierValue,
        value: CompileValue,
    ) -> Result<(), CompileError> {
        self.push(value)?;
        self.get_constant(Value::Identifier(identifier))?;
        self.emit(Instruction::SetGlobal);
        Ok(())
    }

    pub fn compile_declare_macro(
        &mut self,
        identifier: IdentifierValue,
        value: CompileValue,
    ) -> Result<(), CompileError> {
        self.push(value)?;
        self.get_constant(Value::Identifier(identifier))?;
        self.emit(Instruction::DeclareMacro);
        Ok(())
    }

    pub fn compile_assign(
        &mut self,
        identifier: IdentifierValue,
        value: CompileValue,
    ) -> Result<(), CompileError> {
        self.push(value)?;

        if let Some(local) = self.resolve_local(identifier.as_ref())? {
            self.emit(Instruction::SetLocal);
            self.emit(local);
        } else if let Some(upvalue) = self.resolve_upvalue(identifier.as_ref())? {
            self.emit(Instruction::SetUpvalue);
            self.emit(upvalue);
        } else {
            self.get_constant(Value::Identifier(identifier))?;
            self.emit(Instruction::SetGlobal);
        }
        Ok(())
    }

    fn get_constant(&mut self, value: Value) -> Result<(), CompileError> {
        let index = self.constant(value)?;
        self.emit(Instruction::PushConstant);
        self.emit(index);
        Ok(())
    }

    pub fn push(&mut self, value: CompileValue) -> Result<(), CompileError> {
        match value {
            CompileValue::Nil => {
                self.emit(Instruction::PushNil);
            }
            CompileValue::String(value) => {
                self.get_constant(Value::String(value))?;
            }
            CompileValue::Quote(value) => {
                self.get_constant(value.as_ref().clone())?;
            }
            CompileValue::Value(value) => {
                self.get_constant(value)?;
            }
            CompileValue::Integer(value) => {
                if let Ok(immediate) = ImmediateInteger::try_from(value) {
                    self.emit(Instruction::PushInteger);
                    self.emit(immediate);
                } else {
                    self.get_constant(Value::Number(value))?;
                }
            }
            CompileValue::Boolean(BooleanValue(value)) => match value {
                true => self.emit(Instruction::PushTrue),
                false => self.emit(Instruction::PushFalse),
            },
            CompileValue::Identifier(name) => {
                if let Some(local) = self.resolve_local(name.as_ref())? {
                    self.emit(Instruction::GetLocal);
                    self.emit(local);
                } else if let Some(upvalue) = self.resolve_upvalue(name.as_ref())? {
                    self.emit(Instruction::GetUpvalue);
                    self.emit(upvalue);
                } else {
                    self.get_constant(Value::Identifier(name))?;
                    self.emit(Instruction::GetGlobal);
                }
            }
            CompileValue::LocalFunction(index) => {
                let function = self.function_blocks[index].to_bytecode()?;
                self.get_constant(Value::Function(function.clone()))?;
                if !function.upvalues.is_empty() {
                    self.emit(Instruction::MakeClosure);
                }
            }
            CompileValue::Stack => (),
        }
        Ok(())
    }

    pub fn discard(&mut self, value: CompileValue) {
        if matches!(value, CompileValue::Stack) {
            self.emit(Instruction::Drop);
        }
    }

    pub fn push_scope(&mut self) {
        self.scope_depth += 1;
        if self.options.trace_compile {
            eprintln!("COMPILE: push_scope({})", self.scope_depth);
        }
    }

    pub fn pop_scope(&mut self, out_value: Option<CompileValue>) -> Result<(), CompileError> {
        if self.options.trace_compile {
            eprintln!("COMPILE: pop_scope({})", self.scope_depth);
        }
        if self.scope_depth == 0 {
            panic!("Cannot pop out of root scope");
        }
        self.scope_depth -= 1;
        let Some(out_value) = out_value else { todo!() };
        self.push(out_value)?;
        self.emit(Instruction::SetTemp);
        for i in (0..self.locals.len()).rev() {
            if self.locals[i].depth <= self.scope_depth as isize {
                break;
            }
            if self.locals[i].is_captured {
                self.emit(Instruction::CloseUpvalue);
            } else {
                self.emit(Instruction::Drop);
            }
            self.locals.pop();
        }
        self.emit(Instruction::GetTemp);
        Ok(())
    }

    fn resolve_upvalue_inner(
        &mut self,
        at: Option<usize>,
        name: &str,
    ) -> Result<Option<LocalId>, CompileError> {
        let Some(at) = at else {
            return Ok(None);
        };
        let block = &mut self.function_blocks[at];

        if let Some(local) = block.resolve_local(name)? {
            block.locals[usize::from(local)].is_captured = true;

            return self
                .add_upvalue(UpvalueDef {
                    index: local,
                    is_local: true,
                })
                .map(Some);
        }

        let parent = block.parent;
        if let Some(upvalue) = self.resolve_upvalue_inner(parent, name)? {
            return self
                .add_upvalue(UpvalueDef {
                    index: upvalue,
                    is_local: false,
                })
                .map(Some);
        }

        Ok(None)
    }

    fn resolve_upvalue(&mut self, name: &str) -> Result<Option<LocalId>, CompileError> {
        self.resolve_upvalue_inner(self.parent, name)
    }
}

impl Deref for CompileContext {
    type Target = FunctionBlock;

    fn deref(&self) -> &Self::Target {
        &self.function_blocks[self.current]
    }
}
impl DerefMut for CompileContext {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.function_blocks[self.current]
    }
}

impl FunctionBlock {
    fn new(
        parent: Option<usize>,
        identifier: Option<IdentifierValue>,
        docstring: Option<StringValue>,
        signature: &FunctionSignature,
    ) -> Self {
        let mut block = Self {
            parent,
            identifier,
            docstring,
            signature: signature.clone(),
            constants: vec![],
            locals: vec![],
            upvalues: vec![],
            scope_depth: 0,
            instructions: vec![],
            labels: vec![],
            loop_stack: vec![],
        };
        for required in signature.required_arguments.iter() {
            block
                .add_local(required.clone(), Some(-100))
                .expect("couldn't add an argument");
        }
        for optional in signature.optional_arguments.iter() {
            block
                .add_local(optional.clone(), Some(-100))
                .expect("couldn't add an argument");
        }
        block
    }

    fn root(identifier: Option<IdentifierValue>) -> Self {
        Self::new(
            None,
            identifier,
            None,
            &FunctionSignature {
                required_arguments: vec![],
                optional_arguments: vec![],
                rest_argument: None,
            },
        )
    }

    pub fn enter_loop(&mut self, label_entry: usize, label_exit: usize) {
        self.loop_stack.push(LoopStackEntry {
            label_entry,
            label_exit,
        });
    }

    pub fn leave_loop(&mut self) {
        let v = self.loop_stack.pop();
        if v.is_none() {
            panic!("Compiler error: leave_loop() called outside any enclosing loop");
        }
    }

    pub fn emit_break(&mut self) -> Result<(), CompileError> {
        let Some(entry) = self.loop_stack.last() else {
            return Err(CompileError::BreakOutsideOfLoop);
        };
        self.emit(Emitted::Jump(entry.label_exit));
        Ok(())
    }

    pub fn emit_continue(&mut self) -> Result<(), CompileError> {
        let Some(entry) = self.loop_stack.last() else {
            return Err(CompileError::ContinueOutsideOfLoop);
        };
        self.emit(Emitted::Jump(entry.label_entry));
        Ok(())
    }

    pub fn to_bytecode(&self) -> Result<Rc<BytecodeFunction>, CompileError> {
        if !self.loop_stack.is_empty() {
            panic!("Compiler error: loop stack not empty on function completion");
        }

        let mut instructions = vec![];
        let mut resolved_labels = HashMap::new();
        let mut patch_branches = HashMap::new();

        for &offset in self.labels.iter() {
            resolved_labels.insert(offset, None);
        }

        // Emit all code, with placeholders for branch targets
        for (offset, emitted) in self.instructions.iter().enumerate() {
            let ip = instructions.len();
            // Resolve real label address
            if let Some(real) = resolved_labels.get_mut(&offset) {
                assert!(real.is_none());
                *real = Some(ip);
            }
            match emitted {
                &Emitted::Jump(label) => {
                    instructions.push(Instruction::Jump.into());
                    patch_branches.insert(instructions.len(), label);
                    instructions.push(0xFF);
                }
                &Emitted::Branch(label) => {
                    instructions.push(Instruction::Branch.into());
                    patch_branches.insert(instructions.len(), label);
                    instructions.push(0xFF);
                }
                Emitted::LocalId(local) => {
                    instructions.extend_from_slice(&local.to_bytes());
                }
                Emitted::ConstantId(index) => {
                    instructions.extend_from_slice(&index.to_bytes());
                }
                Emitted::ArgumentCount(count) => {
                    instructions.extend_from_slice(&count.to_bytes());
                }
                Emitted::ImmediateInteger(value) => {
                    instructions.extend_from_slice(&value.to_bytes())
                }
                &Emitted::Bool(value) => instructions.push(value as u8),
                &Emitted::Instruction(instruction) => instructions.push(u8::from(instruction)),
            }
        }

        // Patch branch targets
        for (position, label) in patch_branches {
            let label_input = self.labels.get(label).copied().unwrap();
            let branch_target = resolved_labels.get(&label_input).copied().unwrap().unwrap();
            let branch_source = position + 1;
            let branch_offset = branch_target
                .checked_signed_diff(branch_source)
                .and_then(|diff| i8::try_from(diff).ok())
                .ok_or_else(|| CompileError::CannotEmitBranch(branch_source, branch_target))?;
            instructions[position] = branch_offset as u8;
        }

        let min_arity = self.signature.min_arity();
        let max_arity = self.signature.max_arity();

        Ok(Rc::new(BytecodeFunction {
            identifier: self.identifier.clone(),
            instructions: instructions.into(),
            docstring: self.docstring.clone(),
            constants: self.constants.iter().cloned().collect(),
            upvalues: self.upvalues.iter().copied().collect(),
            min_arity,
            max_arity,
        }))
    }

    pub fn new_label(&mut self) -> usize {
        let index = self.labels.len();
        self.labels.push(usize::MAX);
        index
    }

    pub fn adjust_label(&mut self, label: usize) {
        let address = self.instructions.len();
        self.labels[label] = address;
    }

    pub fn emit<I: Into<Emitted>>(&mut self, insn: I) {
        let emitted = insn.into();
        self.instructions.push(emitted);
    }

    pub fn constant(&mut self, value: Value) -> Result<ConstantId, CompileError> {
        if let Some(index) = self.constants.iter().position(|v| v == &value) {
            return Ok(index.try_into().unwrap());
        }
        let index = ConstantId::try_from(self.constants.len())
            .map_err(|_| CompileError::TooManyConstants)?;
        self.constants.push(value);
        Ok(index)
    }

    fn add_local(
        &mut self,
        name: IdentifierValue,
        depth: Option<isize>,
    ) -> Result<LocalId, CompileError> {
        let index =
            LocalId::try_from(self.locals.len()).map_err(|_| CompileError::TooManyLocals)?;
        self.locals.push(Local {
            name,
            is_captured: false,
            depth: depth.unwrap_or(-1),
        });
        Ok(index)
    }

    fn add_upvalue(&mut self, upvalue: UpvalueDef) -> Result<LocalId, CompileError> {
        if let Some(index) = self.upvalues.iter().position(|uv| upvalue == *uv) {
            return Ok(index.try_into().unwrap());
        }
        // Add new
        let index =
            LocalId::try_from(self.upvalues.len()).map_err(|_| CompileError::TooManyLocals)?;
        self.upvalues.push(upvalue);
        Ok(index)
    }

    // Binding management

    pub fn init_local(&mut self, index: LocalId) {
        self.locals[usize::from(index)].depth = self.scope_depth as isize;
    }

    pub fn bind_local(&mut self, name: IdentifierValue) -> Result<LocalId, CompileError> {
        if self.scope_depth == 0 {
            todo!();
        }
        let mut i = self.locals.len();
        while i > 0 {
            i -= 1;
            let local = &self.locals[i];
            if local.depth != -1 && local.depth < self.scope_depth as isize {
                break;
            }
            if local.name == name {
                todo!("Same name, same scope");
            }
        }
        self.add_local(name, None)
    }

    fn resolve_local(&mut self, name: &str) -> Result<Option<LocalId>, CompileError> {
        let Some(index) = self
            .locals
            .iter()
            .rposition(|local| local.name.as_ref() == name)
        else {
            return Ok(None);
        };
        let local = &self.locals[index];
        if local.depth == -1 {
            return Err(CompileError::UndefinedLocalReference(name.into()));
        }

        Ok(Some(index.try_into().unwrap()))
    }
}

#[cfg(test)]
pub fn test_compile(
    expression: &Expression,
) -> Result<(CompileContext, CompileValue), CompileError> {
    let mut cx = CompileContext::new(Default::default(), None);
    let value = expression.compile(&mut cx)?;
    Ok((cx, value))
}