luau-bytecode 0.732.0

Luau bytecode model, builder, serializer, and dumper
Documentation
use luau_common::{BString, ByteSlice};

use super::support::{BuilderClosureNames, BytecodeBuilderScratch};
use super::*;
use crate::dump::{BytecodeDebugRemark, dump_constants, dump_function};
use crate::opcodes::BytecodeTypeTag;
use std::io::Write;

impl<'src> BytecodeBuilder<'src> {
    pub fn set_dump_flags(&mut self, flags: BytecodeDumpFlags) {
        self.dump_flags = flags;
        self.dump_enabled = true;
    }

    pub fn set_dump_source(&mut self, source: impl AsRef<[u8]>) {
        self.dump_source.clear();

        for line in source.as_ref().split(|byte| *byte == b'\n') {
            self.dump_source
                .push(BString::from(line.strip_suffix(b"\r").unwrap_or(line)));
        }
    }

    pub fn needs_debug_remarks(&self) -> bool {
        self.dump_flags.remarks()
    }

    pub fn add_debug_remark(&mut self, text: impl AsRef<[u8]>) {
        if !self.needs_debug_remarks() {
            return;
        }

        let remark = StoredDebugRemark {
            pc: self.scratch.code.len(),
            line: self.current_line,
            text: BString::from(text.as_ref()),
        };
        self.scratch.debug_remarks.push(remark.clone());
        self.dump_remarks.push((remark.line, remark.text.clone()));
    }

    pub fn dump_function(&self, id: usize) -> BString {
        self.functions[id].dump.clone()
    }

    pub fn dump_everything(&self) -> BString {
        let mut result = Vec::new();

        for (index, function) in self.functions.iter().enumerate() {
            write!(result, "Function {index} (").unwrap();
            if function.dump_name.is_empty() {
                result.extend_from_slice(b"??");
            } else {
                result.extend_from_slice(&function.dump_name);
            }
            result.extend_from_slice(b"):\n");
            result.extend_from_slice(&function.dump);
            result.push(b'\n');
        }

        BString::new(result)
    }

    pub fn dump_source_remarks(&self) -> BString {
        let mut result = Vec::new();
        let mut remarks = self
            .dump_remarks
            .iter()
            .map(|(line, text)| (*line, text.as_bstr()))
            .collect::<Vec<_>>();
        remarks.sort();

        let mut next_remark = 0usize;
        for (index, line) in self.dump_source.iter().enumerate() {
            let line_number = i32::try_from(index + 1).unwrap_or(i32::MAX);
            let indent = line
                .iter()
                .take_while(|byte| matches!(byte, b' ' | b'\t'))
                .count();

            while remarks
                .get(next_remark)
                .is_some_and(|remark| remark.0 == line_number)
            {
                result.extend_from_slice(&line[..indent]);
                result.extend_from_slice(b"-- remark: ");
                result.extend_from_slice(remarks[next_remark].1);
                result.push(b'\n');
                next_remark += 1;

                while next_remark < remarks.len()
                    && remarks[next_remark] == remarks[next_remark - 1]
                {
                    next_remark += 1;
                }
            }

            result.extend_from_slice(line);
            if index + 1 < self.dump_source.len() {
                result.push(b'\n');
            }
        }

        BString::new(result)
    }

    pub fn dump_type_info(&self) -> BString {
        let mut result = Vec::new();

        for (index, function) in self.functions.iter().enumerate() {
            if function.type_info.is_empty() {
                continue;
            }

            debug_assert_eq!(function.type_info[0], BytecodeTypeTag::Function as u8);
            debug_assert!(function.type_info.len() >= 2);
            let num_params = function.type_info[1];
            debug_assert!(usize::from(num_params) + 1 < function.type_info.len());

            write!(result, "{index}: function(").unwrap();
            for parameter in 0..num_params {
                let ty = function.type_info[2 + usize::from(parameter)];
                result.extend_from_slice(&self.type_name(ty));
                if parameter + 1 != num_params {
                    result.extend_from_slice(b", ");
                }
            }
            result.extend_from_slice(b")\n");
        }

        BString::new(result)
    }

    pub fn annotate_instruction(
        &self,
        result: &mut BString,
        function_id: u32,
        instruction_pc: u32,
    ) {
        if !self.dump_flags.code() {
            return;
        }

        let function = &self.functions[function_id as usize];
        let offsets = &function.dump_instruction_offsets;
        let mut next = instruction_pc as usize + 1;
        debug_assert!(next < offsets.len());

        while next < offsets.len() && offsets[next] == -1 {
            next += 1;
        }

        let start = offsets[instruction_pc as usize] as usize;
        let end = offsets[next] as usize;
        result.extend_from_slice(&function.dump[start..end]);
    }

    pub(super) fn dump_current_function(&self, id: usize) -> (BString, Vec<i32>) {
        if !self.dump_flags.code() && !self.dump_flags.constants() {
            return (BString::default(), Vec::new());
        }

        let mut result = Vec::new();
        let mut offsets = Vec::new();
        let function = &self.functions[id];
        let closure_names = BuilderClosureNames {
            functions: &self.functions,
        };

        if self.dump_flags.locals() {
            result.extend_from_slice(&self.dump_locals(function, &self.scratch));
        }

        if self.dump_flags.types() {
            result.extend_from_slice(&self.dump_types(function, &self.scratch));
        }

        if self.dump_flags.constants() {
            let borrowed_function = self
                .scratch
                .borrowed_wire(&self.debug_strings, &self.class_shapes);
            result.extend_from_slice(&dump_constants(&borrowed_function, &closure_names));
        }

        let remarks = if self.dump_flags.remarks() {
            self.scratch
                .debug_remarks
                .iter()
                .map(|remark| BytecodeDebugRemark {
                    pc: remark.pc,
                    text: remark.text.as_bstr(),
                })
                .collect::<Vec<_>>()
        } else {
            Vec::new()
        };

        if self.dump_flags.code() {
            let source_lines = self.dump_flags.source().then_some(
                self.dump_source
                    .iter()
                    .map(|line| line.as_bstr())
                    .collect::<Vec<_>>(),
            );
            let code_start = i32::try_from(result.len()).unwrap_or(i32::MAX);
            let borrowed_function = self
                .scratch
                .borrowed_wire(&self.debug_strings, &self.class_shapes);
            let (dump, mut dump_offsets) = dump_function(
                &borrowed_function,
                source_lines.as_deref(),
                self.dump_flags.lines(),
                &remarks,
                &closure_names,
            );
            for offset in &mut dump_offsets {
                if *offset != -1 {
                    *offset = offset.saturating_add(code_start);
                }
            }
            result.extend_from_slice(&dump);
            offsets = dump_offsets;
        }

        (BString::new(result), offsets)
    }

    fn dump_locals(
        &self,
        _function: &BytecodeBuilderFunction,
        scratch: &BytecodeBuilderScratch<'src>,
    ) -> Vec<u8> {
        let mut result = Vec::new();

        for (index, local) in scratch.local_vars.iter().enumerate() {
            if local.start_pc == local.end_pc {
                let line = scratch.lines[local.start_pc as usize];
                writeln!(
                    result,
                    "local {index}: reg {}, start pc {} line {line}, no live range",
                    local.register, local.start_pc
                )
                .unwrap();
            } else {
                let end_pc = local.end_pc - 1;
                let start_line = scratch.lines[local.start_pc as usize];
                let end_line = scratch.lines[end_pc as usize];
                writeln!(
                    result,
                    "local {index}: reg {}, start pc {} line {start_line}, end pc {end_pc} line {end_line}",
                    local.register, local.start_pc
                )
                .unwrap();
            }
        }

        result
    }

    fn dump_types(
        &self,
        function: &BytecodeBuilderFunction,
        scratch: &BytecodeBuilderScratch<'src>,
    ) -> Vec<u8> {
        let mut result = Vec::new();

        for (index, ty) in function.type_info.iter().copied().enumerate().skip(2) {
            write!(result, "R{}: ", index - 2).unwrap();
            result.extend_from_slice(&self.type_name(ty));
            result.extend_from_slice(b" [argument]\n");
        }

        for (index, ty) in scratch.upvalue_types.iter().copied().enumerate() {
            write!(result, "U{index}: ").unwrap();
            result.extend_from_slice(&self.type_name(ty));
            result.push(b'\n');
        }

        for local in &scratch.local_types {
            write!(result, "R{}: ", local.register).unwrap();
            result.extend_from_slice(&self.type_name(local.ty));
            writeln!(result, " from {} to {}", local.start_pc, local.end_pc).unwrap();
        }

        result
    }

    fn type_name(&self, ty: u8) -> Vec<u8> {
        let optional = ty & BytecodeTypeTag::OptionalBit as u8 != 0;
        let tag = ty & !(BytecodeTypeTag::OptionalBit as u8);
        if (BytecodeTypeTag::TaggedUserdataBase as u8..BytecodeTypeTag::TaggedUserdataEnd as u8)
            .contains(&tag)
        {
            let index = usize::from(tag - BytecodeTypeTag::TaggedUserdataBase as u8);
            if let Some(userdata) = self.userdata_types.get(index) {
                let mut result = userdata.name.as_bytes().to_vec();
                if optional {
                    result.push(b'?');
                }
                return result;
            }
        }

        let name = match tag {
            tag if tag == BytecodeTypeTag::Nil as u8 => "nil",
            tag if tag == BytecodeTypeTag::Boolean as u8 => "boolean",
            tag if tag == BytecodeTypeTag::Number as u8 => "number",
            tag if tag == BytecodeTypeTag::Integer as u8 => "integer",
            tag if tag == BytecodeTypeTag::String as u8 => "string",
            tag if tag == BytecodeTypeTag::Table as u8 => "table",
            tag if tag == BytecodeTypeTag::Function as u8 => "function",
            tag if tag == BytecodeTypeTag::Thread as u8 => "thread",
            tag if tag == BytecodeTypeTag::Userdata as u8 => "userdata",
            tag if tag == BytecodeTypeTag::Vector as u8 => "vector",
            tag if tag == BytecodeTypeTag::Buffer as u8 => "buffer",
            tag if tag == BytecodeTypeTag::Any as u8 => "any",
            _ => return format!("unknown({tag})").into_bytes(),
        };
        let mut result = name.as_bytes().to_vec();
        if optional {
            result.push(b'?');
        }
        result
    }
}