lysp/src/vm/value/function.rs

use std::fmt;

use crate::{
    compile::UpvalueDef,
    vm::{
        Value,
        instruction::{ConstantId, ImmediateInteger, Instruction},
        value::{IdentifierValue, StringValue},
    },
};

#[derive(Debug, PartialEq)]
pub struct BytecodeFunction {
    pub identifier: Option<IdentifierValue>,
    pub instructions: Box<[u8]>,
    pub constants: Box<[Value]>,
    pub upvalues: Box<[UpvalueDef]>,
    pub arity: usize,
}

enum TraceArgument {
    Constant,
    ImmediateInteger,
    BranchTarget,
    Byte,
    None,
}

impl fmt::Display for BytecodeFunction {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "<function {}@{:p}>", self.name(), self)
    }
}

impl BytecodeFunction {
    pub fn name(&self) -> &str {
        match self.identifier.as_ref() {
            Some(name) => name.as_ref(),
            None => "unnamed",
        }
    }

    pub fn docstring(&self) -> &StringValue {
        todo!()
    }

    fn trace_immediate_integer_at(&self, address: usize) -> Option<ImmediateInteger> {
        let Some(b0) = self.instructions.get(address).copied() else {
            eprint!(" <??> <??>");
            return None;
        };
        eprint!(" {b0:02x}");
        let Some(b1) = self.instructions.get(address + 1).copied() else {
            eprint!(" <??>");
            return None;
        };
        eprint!(" {b1:02x}");
        Some(ImmediateInteger::from(i16::from_le_bytes([b0, b1])))
    }

    fn trace_constant_at(&self, address: usize) -> Option<Option<&Value>> {
        let Some(b0) = self.instructions.get(address).copied() else {
            eprint!(" <??> <??>");
            return None;
        };
        eprint!(" {b0:02x}");
        let Some(b1) = self.instructions.get(address + 1).copied() else {
            eprint!(" <??>");
            return None;
        };
        eprint!(" {b1:02x}");
        let id = usize::from(ConstantId::from(u16::from_le_bytes([b0, b1])));
        Some(self.constants.get(id))
    }

    fn trace_byte_at(&self, address: usize) -> Option<u8> {
        let Some(b0) = self.instructions.get(address).copied() else {
            eprint!(" <??>");
            return None;
        };
        eprint!(" {b0:02x}");
        Some(b0)
    }

    pub fn disassemble(&self, address: usize, before: usize, after: usize, arrow: bool) {
        let start = address.saturating_sub(before);
        let end = (address + after + 1).min(self.instructions.len());

        let mut position = start;
        while position < end {
            let arrow = if arrow && position == address {
                "==>"
            } else {
                "   "
            };
            eprint!("{self:p}:{position:<4}{arrow} ");

            let opcode = self.instructions[position];

            eprint!("{opcode:02x}");

            let Ok(instruction) = Instruction::try_from(opcode) else {
                eprintln!(" <UNKNOWN>");
                position += 1;
                continue;
            };

            let argument = match instruction {
                // Stack
                Instruction::PushNil => TraceArgument::None,
                Instruction::PushTrue | Instruction::PushFalse => TraceArgument::None,
                Instruction::PushInteger => TraceArgument::ImmediateInteger,
                Instruction::PushConstant => TraceArgument::Constant,
                Instruction::SetTemp | Instruction::GetTemp => TraceArgument::None,
                Instruction::SetLocal | Instruction::GetLocal => TraceArgument::Byte,
                Instruction::SetGlobal | Instruction::GetGlobal => TraceArgument::None,
                Instruction::SetUpvalue | Instruction::GetUpvalue => TraceArgument::Byte,
                Instruction::Drop => TraceArgument::None,
                Instruction::DeclareGlobal => TraceArgument::None,
                Instruction::DeclareMacro => TraceArgument::None,
                // Arithmetic
                Instruction::Ne
                | Instruction::Gt
                | Instruction::Lt
                | Instruction::Eq
                | Instruction::Ge
                | Instruction::Le
                | Instruction::Add
                | Instruction::Sub
                | Instruction::Mul
                | Instruction::Div
                | Instruction::Mod
                | Instruction::Negate
                | Instruction::Not => TraceArgument::Byte,
                // Function
                Instruction::Call => TraceArgument::Byte,
                Instruction::Return => TraceArgument::None,
                Instruction::MakeClosure => TraceArgument::None,
                Instruction::CloseUpvalue => TraceArgument::Byte,
                // Branch
                Instruction::Branch => TraceArgument::BranchTarget,
                Instruction::Jump => TraceArgument::BranchTarget,
            };

            position += 1;

            match argument {
                TraceArgument::Byte => {
                    let Some(byte) = self.trace_byte_at(position) else {
                        eprintln!("\t\t{instruction} <invalid>");
                        break;
                    };
                    position += 1;
                    eprintln!("\t\t{instruction} {byte}");
                }
                TraceArgument::Constant => {
                    let Some(constant) = self.trace_constant_at(position) else {
                        eprintln!("\t\t{instruction} <invalid>");
                        break;
                    };
                    position += 2;
                    let Some(constant) = constant else {
                        eprintln!("\t\t{instruction} <invalid>");
                        continue;
                    };
                    eprintln!("\t\t{instruction} {constant}");
                }
                TraceArgument::ImmediateInteger => {
                    let Some(value) = self.trace_immediate_integer_at(position) else {
                        eprintln!("\t\t{instruction} <invalid>");
                        break;
                    };
                    position += 2;
                    eprintln!("\t\t{instruction} {}", value.sign_extend_i64());
                }
                TraceArgument::None => {
                    eprintln!("\t\t{instruction}");
                }
                TraceArgument::BranchTarget => {
                    let Some(byte) = self.trace_byte_at(position) else {
                        eprintln!("\t\t{instruction} <invalid>");
                        break;
                    };
                    position += 1;
                    eprintln!(
                        "\t\t{instruction} {:+} (-> {})",
                        byte as i8,
                        position.saturating_add_signed(isize::from(byte as i8))
                    );
                }
            }
        }
    }
}