luau-bytecode 0.732.0

Luau bytecode model, builder, serializer, and dumper
Documentation
use super::support::{BytecodeBuilderLocal, BytecodeBuilderScratch, Jump};
use super::*;
use crate::model::{Instruction, InstructionWord, Register};
use crate::opcodes::Opcode;

impl<'src> BytecodeBuilder<'src> {
    pub fn emit_abc(&mut self, opcode: Opcode, a: u8, b: u8, c: u8) {
        self.emit(Instruction::abc(opcode, a, b, c));
    }

    pub fn set_debug_function_name(&mut self, name: impl Into<BytecodeStringRef<'src>>) {
        let name = name.into();
        let index = self.add_string_table_entry(&name);
        let dump_enabled = self.dump_enabled;
        let function = self.current_function_meta();
        function.debug_name = Some(index);
        if dump_enabled {
            function.dump_name.clear();
            function.dump_name.extend_from_slice(name.as_bytes());
        }
    }

    pub fn set_debug_function_line_defined(&mut self, line: i32) {
        self.current_function_meta().line_defined = line;
    }

    pub fn add_flags(&mut self, flags: u8) {
        self.current_function_meta().flags |= flags;
    }

    pub fn push_open_debug_local(
        &mut self,
        name: impl Into<BytecodeStringRef<'src>>,
        register: Register,
    ) {
        let start_pc = self.current_function().code.len() as u32;
        self.push_debug_local(name, register, start_pc, u32::MAX);
    }

    pub fn push_debug_local(
        &mut self,
        name: impl Into<BytecodeStringRef<'src>>,
        register: Register,
        start_pc: u32,
        end_pc: u32,
    ) {
        let name = name.into();
        let name_ref = self.add_string_table_entry(&name);
        self.current_function()
            .local_vars
            .push(BytecodeBuilderLocal {
                name: name_ref,
                start_pc,
                end_pc,
                register,
            });
    }

    pub fn close_debug_locals_from(&mut self, register: Register) {
        let end_pc = self.current_function().code.len() as u32;
        for local in self
            .current_function()
            .local_vars
            .iter_mut()
            .filter(|local| local.register >= register && local.end_pc == u32::MAX)
        {
            local.end_pc = end_pc;
        }
    }

    pub fn push_debug_upvalue(&mut self, name: impl Into<BytecodeStringRef<'src>>) {
        let name = name.into();
        let name_ref = self.add_string_table_entry(&name);
        self.current_function().upvalues.push(name_ref);
    }

    pub fn emit_ad(&mut self, opcode: Opcode, a: u8, d: i16) {
        self.emit(Instruction::ad(opcode, a, d));
    }

    pub fn emit_e(&mut self, opcode: Opcode, e: i32) {
        self.emit(Instruction::ae(opcode, e));
    }

    pub fn set_debug_line(&mut self, line: usize) {
        self.current_line = line.try_into().unwrap_or(i32::MAX);
    }

    pub fn current_line(&self) -> i32 {
        self.current_line
    }

    pub fn restore_line(&mut self, line: i32) {
        self.current_line = line;
    }

    pub fn emit_aux(&mut self, aux: InstructionWord) {
        let line = self.current_line;
        self.current_function().code.push(Instruction::new(aux));
        self.current_function().lines.push(line);
    }

    pub fn undo_emit(&mut self, opcode: Opcode) {
        let instruction = self
            .current_function()
            .code
            .pop()
            .expect("BytecodeBuilder::undo_emit requires an emitted instruction");
        debug_assert_eq!(unsafe { instruction.opcode_unchecked() }, opcode);
        self.current_function().lines.pop();
    }

    pub fn emit_label(&mut self) -> usize {
        self.current_function().code.len()
    }

    pub fn patch_jump_d(&mut self, jump_label: usize, target_label: usize) -> bool {
        let offset = target_label as isize - jump_label as isize - 1;
        let jump_instruction = self
            .current_function_ref()
            .code
            .get(jump_label)
            .copied()
            .expect("jump label must point at an emitted instruction");
        debug_assert!(unsafe { jump_instruction.opcode_unchecked() }.is_jump_d());
        debug_assert_eq!(jump_instruction.d(), 0);
        debug_assert!(target_label <= self.current_function_ref().code.len());

        if let Ok(offset) = i16::try_from(offset) {
            let instruction = self
                .current_function()
                .code
                .get_mut(jump_label)
                .expect("jump label must point at an emitted instruction");
            *instruction = instruction.with_d(offset);
        } else if offset.unsigned_abs() < (1 << 23) {
            self.current_function().has_long_jumps = true;
        } else {
            return false;
        }

        self.current_function().jumps.push(Jump {
            source: jump_label,
            target: target_label,
        });
        true
    }

    pub fn patch_jump_e(&mut self, jump_label: usize, target_label: usize) -> bool {
        let offset = target_label as isize - jump_label as isize - 1;
        let jump_instruction = self
            .current_function_ref()
            .code
            .get(jump_label)
            .copied()
            .expect("jump label must point at an emitted instruction");
        debug_assert_eq!(
            unsafe { jump_instruction.opcode_unchecked() },
            Opcode::JumpX
        );
        debug_assert_eq!(jump_instruction.e(), 0);
        debug_assert!(target_label <= self.current_function_ref().code.len());

        if !(-(1 << 23)..(1 << 23)).contains(&offset) {
            return false;
        }

        let instruction = self
            .current_function()
            .code
            .get_mut(jump_label)
            .expect("jump label must point at an emitted instruction");
        *instruction = Instruction::ae(Opcode::JumpX, offset as i32);
        true
    }

    pub fn patch_skip_c(&mut self, jump_label: usize, target_label: usize) -> bool {
        let offset = target_label as isize - jump_label as isize - 1;
        let jump_instruction = self
            .current_function_ref()
            .code
            .get(jump_label)
            .copied()
            .expect("jump label must point at an emitted instruction");
        debug_assert!(
            unsafe { jump_instruction.opcode_unchecked() }.is_skip_c()
                || unsafe { jump_instruction.opcode_unchecked() }.is_fast_call()
        );
        debug_assert_eq!(jump_instruction.c(), 0);
        let Ok(offset) = u8::try_from(offset) else {
            return false;
        };

        let instruction = self
            .current_function()
            .code
            .get_mut(jump_label)
            .expect("jump label must point at an emitted instruction");
        *instruction = instruction.with_c(offset);
        true
    }

    pub fn patch_aux(&mut self, target_aux: usize, value: i32) {
        let instruction = self
            .current_function()
            .code
            .get_mut(target_aux)
            .expect("aux patch target must point at an emitted instruction");
        *instruction = Instruction::new(value as u32);
    }

    pub fn fold_jumps(&mut self) {
        if self.current_function_ref().has_long_jumps {
            return;
        }

        for jump_index in 0..self.current_function_ref().jumps.len() {
            let jump_label = self.current_function_ref().jumps[jump_index].source;
            let jump_instruction = self.current_function().code[jump_label];
            let mut target_label: usize =
                (jump_label as isize + 1 + isize::from(jump_instruction.d()))
                    .try_into()
                    .expect("jump target must be non-negative");
            let mut target_instruction = self.current_function().code[target_label];

            while unsafe { target_instruction.opcode_unchecked() } == Opcode::Jump
                && target_instruction.d() >= 0
            {
                target_label = (target_label as isize + 1 + isize::from(target_instruction.d()))
                    .try_into()
                    .expect("jump target must be non-negative");
                target_instruction = self.current_function().code[target_label];
            }

            let offset = target_label as isize - jump_label as isize - 1;
            let instruction = self
                .current_function()
                .code
                .get_mut(jump_label)
                .expect("jump label must point at an emitted instruction");

            if unsafe { jump_instruction.opcode_unchecked() } == Opcode::Jump
                && unsafe { target_instruction.opcode_unchecked() } == Opcode::Return
            {
                *instruction = target_instruction;
            } else if let Ok(offset) = i16::try_from(offset) {
                *instruction = instruction.with_d(offset);
            }

            self.current_function().jumps[jump_index].target = target_label;
        }
    }

    pub fn expand_jumps(&mut self) {
        if !self.current_function_ref().has_long_jumps {
            return;
        }

        const MAX_JUMP_DISTANCE_CONSERVATIVE: isize = 32767 / 3;

        let (jumps, old_code, old_lines) = {
            let function = self.current_function();
            function.jumps.sort_by_key(|jump| jump.source);
            (
                std::mem::take(&mut function.jumps),
                std::mem::take(&mut function.code),
                std::mem::take(&mut function.lines),
            )
        };
        let mut remap = vec![0usize; old_code.len()];
        let mut new_code = Vec::with_capacity(old_code.len());
        let mut new_lines = Vec::with_capacity(old_lines.len());

        let mut current_jump = 0usize;
        let mut pending_trampolines = 0usize;

        let mut pc = 0usize;
        while pc < old_code.len() {
            let instruction = old_code[pc];
            if current_jump < jumps.len() && jumps[current_jump].source == pc {
                let offset =
                    jumps[current_jump].target as isize - jumps[current_jump].source as isize - 1;

                if offset.abs() > MAX_JUMP_DISTANCE_CONSERVATIVE {
                    new_code.push(Instruction::ad(Opcode::Jump, 0, 1));
                    new_code.push(Instruction::ae(Opcode::JumpX, 0));
                    new_lines.push(old_lines[pc]);
                    new_lines.push(old_lines[pc]);
                    pending_trampolines += 1;
                }

                current_jump += 1;
            }

            let opcode = unsafe { instruction.opcode_unchecked() };
            for word_pc in pc..pc + opcode.length() {
                remap[word_pc] = new_code.len();
                new_code.push(old_code[word_pc]);
                new_lines.push(old_lines[word_pc]);
            }
            pc += opcode.length();
        }

        for jump in &jumps {
            let offset = jump.target as isize - jump.source as isize - 1;
            let new_offset = remap[jump.target] as isize - remap[jump.source] as isize - 1;

            if offset.abs() > MAX_JUMP_DISTANCE_CONSERVATIVE {
                let trampoline = remap[jump.source] - 1;
                new_code[trampoline] = Instruction::ae(Opcode::JumpX, (new_offset + 1) as i32);
                new_code[remap[jump.source]] = new_code[remap[jump.source]].with_d(-2);
                pending_trampolines -= 1;
            } else {
                let new_offset =
                    i16::try_from(new_offset).expect("expanded jump offset must fit i16");
                new_code[remap[jump.source]] = new_code[remap[jump.source]].with_d(new_offset);
            }
        }

        debug_assert_eq!(pending_trampolines, 0);

        let function = self.current_function();
        function.jumps = jumps;
        function.code = new_code;
        function.lines = new_lines;
        for local in &mut function.local_vars {
            local.end_pc = if local.start_pc != local.end_pc {
                remap[local.end_pc as usize - 1] as u32 + 1
            } else {
                remap[local.end_pc as usize] as u32
            };
            local.start_pc = remap[local.start_pc as usize] as u32;
        }
        for local in &mut function.local_types {
            local.end_pc = if local.start_pc != local.end_pc {
                remap[local.end_pc as usize - 1] as u32 + 1
            } else {
                remap[local.end_pc as usize] as u32
            };
            local.start_pc = remap[local.start_pc as usize] as u32;
        }
    }

    pub fn get_instruction_count(&self) -> usize {
        self.current_function_ref().code.len()
    }

    pub fn get_total_instruction_count(&self) -> usize {
        self.total_instruction_count
    }

    pub fn get_debug_pc(&self) -> u32 {
        self.get_instruction_count() as u32
    }

    pub(super) fn current_function(&mut self) -> &mut BytecodeBuilderScratch<'src> {
        &mut self.scratch
    }

    pub(super) fn current_function_ref(&self) -> &BytecodeBuilderScratch<'src> {
        &self.scratch
    }

    pub(super) fn current_function_meta(&mut self) -> &mut BytecodeBuilderFunction {
        let id = self.current_function_id();
        &mut self.functions[id]
    }

    pub(super) fn current_function_id(&self) -> usize {
        self.current_function
            .expect("bytecode emission requires an active function")
    }

    pub(super) fn emit(&mut self, instruction: Instruction) {
        let line = self.current_line;
        let proto = self.current_function();
        proto.code.push(instruction);
        proto.lines.push(line);
    }
}