luau-bytecode 0.732.0

mod debug;
mod emit;
mod serialize;
mod support;

use crate::model::{BytecodeClass, BytecodeUserdataType, InstructionWord};
use luau_common::{BString, DenseHashMap};

use support::{
    BytecodeBuilderFunction, BytecodeBuilderScratch, BytecodeStringHasher, StoredDebugRemark,
};
pub use support::{BytecodeDumpFlags, BytecodeEncoder};

const MAX_CONSTANT_COUNT: usize = 1 << 23;

pub struct BytecodeBuilder<'src> {
    functions: Vec<BytecodeBuilderFunction>,
    scratch: BytecodeBuilderScratch<'src>,
    class_shapes: Vec<BytecodeClass>,
    userdata_types: Vec<BytecodeUserdataType>,
    string_index: DenseHashMap<BytecodeStringRef<'src>, u32, BytecodeStringHasher>,
    debug_strings: Vec<BytecodeStringRef<'src>>,
    bytecode: Vec<u8>,
    encoder: Option<Box<dyn BytecodeEncoder>>,
    current_function: Option<usize>,
    main: Option<usize>,
    total_instruction_count: usize,
    current_line: i32,
    dump_flags: BytecodeDumpFlags,
    dump_enabled: bool,
    dump_source: Vec<BString>,
    dump_remarks: Vec<(i32, BString)>,
}

#[derive(Debug, Clone, Copy)]
pub struct BytecodeStringRef<'src> {
    data: *const u8,
    length: usize,
    _marker: std::marker::PhantomData<&'src [u8]>,
}

impl<'src> BytecodeStringRef<'src> {
    pub const fn empty() -> Self {
        Self {
            data: std::ptr::null(),
            length: 0,
            _marker: std::marker::PhantomData,
        }
    }

    pub fn as_bytes(&self) -> &'src [u8] {
        if self.data.is_null() {
            &[]
        } else {
            unsafe { std::slice::from_raw_parts(self.data, self.length) }
        }
    }
}

impl PartialEq for BytecodeStringRef<'_> {
    fn eq(&self, other: &Self) -> bool {
        if !self.data.is_null() && !other.data.is_null() {
            self.length == other.length && self.as_bytes() == other.as_bytes()
        } else {
            self.data == other.data
        }
    }
}

impl Eq for BytecodeStringRef<'_> {}

impl<'src> From<&'src [u8]> for BytecodeStringRef<'src> {
    fn from(value: &'src [u8]) -> Self {
        Self {
            data: value.as_ptr(),
            length: value.len(),
            _marker: std::marker::PhantomData,
        }
    }
}

impl<'src, const N: usize> From<&'src [u8; N]> for BytecodeStringRef<'src> {
    fn from(value: &'src [u8; N]) -> Self {
        Self::from(value.as_slice())
    }
}

impl<'src> Default for BytecodeBuilder<'src> {
    fn default() -> Self {
        Self {
            functions: Vec::new(),
            scratch: BytecodeBuilderScratch::default(),
            class_shapes: Vec::new(),
            userdata_types: Vec::new(),
            string_index: DenseHashMap::new(BytecodeStringRef::empty()),
            debug_strings: Vec::new(),
            bytecode: Vec::new(),
            encoder: None,
            current_function: None,
            main: None,
            total_instruction_count: 0,
            current_line: 0,
            dump_flags: BytecodeDumpFlags::default(),
            dump_enabled: false,
            dump_source: Vec::new(),
            dump_remarks: Vec::new(),
        }
    }
}

impl<'src> BytecodeBuilder<'src> {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_encoder(encoder: impl BytecodeEncoder + 'static) -> Self {
        Self {
            encoder: Some(Box::new(encoder)),
            ..Self::default()
        }
    }

    pub fn begin_function(&mut self, num_params: u8, is_vararg: bool) -> usize {
        debug_assert!(
            self.current_function.is_none(),
            "BytecodeBuilder::begin_function requires no active function"
        );

        let id = self.functions.len();
        self.functions
            .push(BytecodeBuilderFunction::new(num_params, is_vararg));
        self.current_function = Some(id);
        self.current_line = 0;
        id
    }

    pub fn function_count(&self) -> usize {
        self.functions.len()
    }

    pub fn clear_string_table(&mut self) {
        self.string_index.clear();
    }

    pub fn end_function(
        &mut self,
        max_stack_size: u8,
        upvalue_count: u8,
        flags: u8,
        cost: u64,
    ) -> usize {
        let id = self
            .current_function
            .take()
            .expect("BytecodeBuilder::end_function requires an active function");
        self.functions[id].max_stack_size = max_stack_size;
        self.functions[id].upvalue_count = upvalue_count;
        self.functions[id].flags |= flags;
        self.functions[id].cost = cost;
        #[cfg(debug_assertions)]
        self.scratch.validate(&self.functions[id], &self.functions);
        let (dump, dump_instruction_offsets) = self.dump_current_function(id);
        self.functions[id].dump = dump;
        self.functions[id].dump_instruction_offsets = dump_instruction_offsets;
        if let Some(encoder) = self.encoder.as_deref() {
            let code = &mut self.scratch.code;
            // Safety: `Instruction` is `repr(transparent)` over `InstructionWord`, so the live
            // instruction buffer has the exact layout the encoder expects.
            let words = unsafe {
                std::slice::from_raw_parts_mut(
                    code.as_mut_ptr().cast::<InstructionWord>(),
                    code.len(),
                )
            };
            encoder.encode(words);
        }
        self.functions[id].data = self.function_data(id);
        self.total_instruction_count += self.scratch.code.len();
        self.scratch.clear();
        id
    }

    pub fn add_child_function(&mut self, id: u32) -> Option<i16> {
        self.current_function
            .expect("BytecodeBuilder::add_child_function requires an active function");

        if let Some(index) = self.scratch.child_function_map.get(&id) {
            return Some(*index);
        }

        let index = i16::try_from(self.scratch.child_functions.len()).ok()?;
        self.scratch.child_functions.push(id);
        let (_, fresh) = self.scratch.child_function_map.insert(id, index);
        debug_assert!(fresh);
        Some(index)
    }

    pub fn set_main_function(&mut self, id: usize) {
        debug_assert!(
            id < self.functions.len(),
            "BytecodeBuilder::set_main_function requires a valid function id"
        );
        self.main = Some(id);
    }
}