luau-bytecode 0.732.0

use luau_common::{BString, ByteSlice};
use std::borrow::Borrow;

use crate::opcodes::{InvalidOpcode, Opcode};

/// VM register index encoded in bytecode.
pub type Register = u8;

/// Proto constant table index encoded in bytecode.
pub type ConstantIndex = i32;

/// Packed table-template constant value encoded in bytecode.
pub type PackedTableValue = i32;

/// Raw bytecode instruction word.
pub type InstructionWord = u32;

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct BytecodeClass {
    pub class_name: ConstantIndex,
    pub property_names: Vec<ConstantIndex>,
    pub method_names: Vec<ConstantIndex>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct BytecodeImportId(u32);

impl BytecodeImportId {
    pub fn new(id0: ConstantIndex) -> Self {
        debug_assert!((0..1024).contains(&id0));
        Self((1 << 30) | ((id0 as u32) << 20))
    }

    pub fn from_two(id0: ConstantIndex, id1: ConstantIndex) -> Self {
        debug_assert!((0..1024).contains(&id0));
        debug_assert!((0..1024).contains(&id1));
        Self((2 << 30) | ((id0 as u32) << 20) | ((id1 as u32) << 10))
    }

    pub fn from_three(id0: ConstantIndex, id1: ConstantIndex, id2: ConstantIndex) -> Self {
        debug_assert!((0..1024).contains(&id0));
        debug_assert!((0..1024).contains(&id1));
        debug_assert!((0..1024).contains(&id2));
        Self((3 << 30) | ((id0 as u32) << 20) | ((id1 as u32) << 10) | id2 as u32)
    }

    pub fn from_raw(value: u32) -> Self {
        Self(value)
    }

    pub fn raw(self) -> u32 {
        self.0
    }

    pub fn components(self) -> Vec<ConstantIndex> {
        let count = self.0 >> 30;
        let mut result = Vec::with_capacity(count as usize);
        if count > 0 {
            result.push(((self.0 >> 20) & 1023) as ConstantIndex);
        }
        if count > 1 {
            result.push(((self.0 >> 10) & 1023) as ConstantIndex);
        }
        if count > 2 {
            result.push((self.0 & 1023) as ConstantIndex);
        }
        result
    }
}

#[derive(Debug, Clone, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct BytecodeString(BString);

impl BytecodeString {
    pub fn as_bytes(&self) -> &[u8] {
        self.0.as_bytes()
    }

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

    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
}

impl Borrow<[u8]> for BytecodeString {
    fn borrow(&self) -> &[u8] {
        self.as_bytes()
    }
}

impl From<&[u8]> for BytecodeString {
    fn from(value: &[u8]) -> Self {
        Self(BString::from(value))
    }
}

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

impl From<Vec<u8>> for BytecodeString {
    fn from(value: Vec<u8>) -> Self {
        Self(BString::new(value))
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ClosureIndex(u32);

impl ClosureIndex {
    pub fn new(index: u32) -> Self {
        Self(index)
    }

    pub fn get(self) -> u32 {
        self.0
    }

    pub fn as_usize(self) -> usize {
        self.0 as usize
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TableShape {
    entries: [TableShapeEntry; Self::MAX_LENGTH],
    length: u8,
    has_constants: bool,
}

impl TableShape {
    pub const MAX_LENGTH: usize = 32;

    pub fn new(entries: Vec<TableShapeEntry>) -> Self {
        debug_assert!(entries.len() <= Self::MAX_LENGTH);
        let has_constants = entries.iter().any(|entry| entry.value.is_some());
        let mut storage = [TableShapeEntry::default(); Self::MAX_LENGTH];
        let length = entries.len();
        storage[..length].copy_from_slice(&entries);
        Self {
            entries: storage,
            length: length as u8,
            has_constants,
        }
    }

    pub fn entries(&self) -> &[TableShapeEntry] {
        &self.entries[..usize::from(self.length)]
    }

    pub fn len(&self) -> usize {
        usize::from(self.length)
    }

    pub fn is_empty(&self) -> bool {
        self.length == 0
    }

    pub fn has_constants(&self) -> bool {
        self.has_constants
    }
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash)]
pub struct TableShapeEntry {
    pub key: ConstantIndex,
    pub value: Option<PackedTableValue>,
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct BytecodeVector {
    values: [f32; 4],
}

impl BytecodeVector {
    pub fn new(x: f32, y: f32, z: f32, w: f32) -> Self {
        Self {
            values: [x, y, z, w],
        }
    }

    pub fn x(self) -> f32 {
        self.values[0]
    }

    pub fn y(self) -> f32 {
        self.values[1]
    }

    pub fn z(self) -> f32 {
        self.values[2]
    }

    pub fn w(self) -> f32 {
        self.values[3]
    }

    pub fn to_bits(self) -> [u32; 4] {
        self.values.map(f32::to_bits)
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct BytecodeVectorDouble {
    values: [f64; 4],
}

impl BytecodeVectorDouble {
    pub fn new(x: f64, y: f64, z: f64, w: f64) -> Self {
        Self {
            values: [x, y, z, w],
        }
    }

    pub fn x(self) -> f64 {
        self.values[0]
    }

    pub fn y(self) -> f64 {
        self.values[1]
    }

    pub fn z(self) -> f64 {
        self.values[2]
    }

    pub fn w(self) -> f64 {
        self.values[3]
    }

    pub fn to_bits(self) -> [u64; 4] {
        self.values.map(f64::to_bits)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BytecodeTypedLocal {
    pub ty: u8,
    pub register: Register,
    pub start_pc: u32,
    pub end_pc: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BytecodeFeedbackSlot {
    pub pc: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BytecodeFeedbackType {
    CallTarget,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct BytecodeUserdataType {
    pub name: BytecodeString,
    pub name_ref: u32,
    pub used: bool,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
pub struct Instruction(InstructionWord);

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
pub struct InstructionAux(InstructionWord);

impl Instruction {
    pub fn new(word: InstructionWord) -> Self {
        Self(word)
    }

    pub fn abc(opcode: Opcode, a: u8, b: u8, c: u8) -> Self {
        Self(opcode.word() | ((a as u32) << 8) | ((b as u32) << 16) | ((c as u32) << 24))
    }

    pub fn ad(opcode: Opcode, a: u8, d: i16) -> Self {
        Self(opcode.word() | ((a as u32) << 8) | (((d as u16) as u32) << 16))
    }

    pub fn ae(opcode: Opcode, e: i32) -> Self {
        Self(opcode.word() | (((e as u32) & 0x00ff_ffff) << 8))
    }

    pub fn word(self) -> InstructionWord {
        self.0
    }

    pub fn with_d(self, d: i16) -> Self {
        Self((self.0 & 0x0000_ffff) | (((d as u16) as u32) << 16))
    }

    pub fn with_c(self, c: u8) -> Self {
        Self((self.0 & 0x00ff_ffff) | ((c as u32) << 24))
    }

    /// Decodes the opcode byte when this word is an instruction header.
    pub fn try_opcode(self) -> Result<Opcode, InvalidOpcode> {
        let opcode = (self.0 & 0xff) as u8;
        Opcode::from_byte(opcode).ok_or(InvalidOpcode::new(opcode))
    }

    /// Decodes the opcode byte without checking that it is in range.
    ///
    /// # Safety
    ///
    /// The word must be an instruction header from bytecode that has already
    /// passed opcode validation. AUX words and VM cache/sentinel words must not
    /// be decoded through this function.
    pub unsafe fn opcode_unchecked(self) -> Opcode {
        unsafe { self.try_opcode().unwrap_unchecked() }
    }

    pub fn a(self) -> u8 {
        ((self.0 >> 8) & 0xff) as u8
    }

    pub fn b(self) -> u8 {
        ((self.0 >> 16) & 0xff) as u8
    }

    pub fn c(self) -> u8 {
        ((self.0 >> 24) & 0xff) as u8
    }

    pub fn d(self) -> i16 {
        (self.0 >> 16) as u16 as i16
    }

    pub fn e(self) -> i32 {
        (self.0 as i32) >> 8
    }

    /// Computes the jump target for a checked instruction header.
    pub fn try_jump_target(self, pc: u32) -> Result<Option<i32>, InvalidOpcode> {
        Ok(self.try_opcode()?.jump_target(self, pc))
    }

    /// Computes the jump target without checking the opcode byte.
    ///
    /// # Safety
    ///
    /// The word must satisfy the safety requirements of [`Self::opcode_unchecked`].
    /// Non-jump opcodes are valid inputs and return `None`.
    pub unsafe fn jump_target_unchecked(self, pc: u32) -> Option<i32> {
        unsafe { self.opcode_unchecked().jump_target(self, pc) }
    }
}

impl InstructionAux {
    pub fn new(word: InstructionWord) -> Self {
        Self(word)
    }

    pub fn word(self) -> InstructionWord {
        self.0
    }

    pub fn a(self) -> u8 {
        (self.0 & 0xff) as u8
    }

    pub fn b(self) -> u8 {
        ((self.0 >> 8) & 0xff) as u8
    }

    pub fn kv(self) -> u32 {
        self.0 & 0x00ff_ffff
    }

    pub fn kb(self) -> bool {
        (self.0 & 1) != 0
    }

    pub fn is_negated(self) -> bool {
        (self.0 >> 31) != 0
    }

    pub fn kv16(self) -> u16 {
        (self.0 & 0xffff) as u16
    }

    pub fn slot(self) -> u16 {
        (self.0 >> 16) as u16
    }
}