use crate::builder::BytecodeBuilder;
use crate::error::BytecodeReadError;
use crate::graph::{
BytecodeBlock, BytecodeBlockId, BytecodeImmediate, BytecodeImmediateId, BytecodeInstruction,
BytecodeInstructionId, BytecodeOperand, BytecodePhi, BytecodePhiId, BytecodeProjection,
BytecodeProjectionId, BytecodeWriteError, InstructionPc, build_function_graph,
encode_function_bytecode,
};
use crate::model::{
BytecodeClass, BytecodeTypedLocal, BytecodeVector, BytecodeVectorDouble, Instruction,
InstructionWord, Register, TableShape, TableShapeEntry,
};
use crate::opcodes::{BytecodeConstantTag, Opcode};
use luau_common::{bytecode_wire, flags};
use std::borrow::Cow;
use std::collections::HashMap;
#[derive(Debug, Clone, PartialEq)]
pub struct BytecodeFunction<'table> {
pub max_stack_size: u8,
pub num_params: u8,
pub upvalue_count: u8,
pub is_vararg: bool,
pub flags: u8,
pub type_info: Vec<u8>,
pub upvalue_types: Vec<u8>,
pub local_types: Vec<BytecodeTypedLocal>,
pub blocks: Vec<BytecodeBlock>,
pub instructions: Vec<BytecodeInstruction>,
pub constants: Vec<BytecodeFunctionConstant<'table>>,
pub immediates: Vec<BytecodeImmediate>,
pub phis: Vec<BytecodePhi>,
pub projections: Vec<BytecodeProjection>,
pub registers: HashMap<BytecodeOperand, Register>,
pub table_shapes: Vec<TableShape>,
pub class_shapes: Vec<BytecodeClass>,
pub entry_block: BytecodeBlockId,
pub exit_block: BytecodeBlockId,
pub pc_to_block: Vec<BytecodeBlockId>,
pub pc_to_instruction: Vec<BytecodeInstructionId>,
pub protos: Vec<u32>,
pub line_defined: u32,
pub debug_name: &'table [u8],
pub lines: Vec<u32>,
pub locals: Vec<BytecodeDebugLocal<'table>>,
pub upvalue_names: Vec<&'table [u8]>,
}
impl<'table> BytecodeFunction<'table> {
pub fn from_function_bytecode<'strings>(
data: &[u8],
strings: &'table BytecodeStringTable<'strings>,
) -> Result<Self, BytecodeReadError> {
BytecodeFunctionReader::new(data, strings).read()
}
pub fn to_function_bytecode(&mut self) -> Result<Vec<u8>, BytecodeWriteError> {
let mut builder = BytecodeBuilder::new();
encode_function_bytecode(&mut builder, self)
}
pub fn entry(&self) -> BytecodeBlockId {
self.entry_block
}
pub fn exit(&self) -> BytecodeBlockId {
self.exit_block
}
pub fn blocks(&self) -> &[BytecodeBlock] {
&self.blocks
}
pub fn block(&self, id: BytecodeBlockId) -> &BytecodeBlock {
&self.blocks[id.index()]
}
pub fn graph_instruction(&self, id: BytecodeInstructionId) -> &BytecodeInstruction {
&self.instructions[id.index()]
}
pub fn block_for_pc(&self, pc: InstructionPc) -> Option<BytecodeBlockId> {
self.pc_to_block.get(pc.index()).copied()
}
pub fn instruction_for_pc(&self, pc: InstructionPc) -> Option<BytecodeInstructionId> {
self.pc_to_instruction.get(pc.index()).copied()
}
pub fn immediate(&self, id: BytecodeImmediateId) -> &BytecodeImmediate {
&self.immediates[id.index()]
}
pub fn phi(&self, id: BytecodePhiId) -> &BytecodePhi {
&self.phis[id.index()]
}
pub fn projection(&self, id: BytecodeProjectionId) -> &BytecodeProjection {
&self.projections[id.index()]
}
}
struct BytecodeFunctionReader<'data, 'table, 'strings> {
data: &'data [u8],
offset: usize,
strings: &'table BytecodeStringTable<'strings>,
table_shapes: Vec<TableShape>,
class_shapes: Vec<BytecodeClass>,
}
impl<'data, 'table, 'strings> BytecodeFunctionReader<'data, 'table, 'strings> {
fn new(data: &'data [u8], strings: &'table BytecodeStringTable<'strings>) -> Self {
Self {
data,
offset: 0,
strings,
table_shapes: Vec::new(),
class_shapes: Vec::new(),
}
}
fn read(mut self) -> Result<BytecodeFunction<'table>, BytecodeReadError> {
let max_stack_size = self.read_u8()?;
let num_params = self.read_u8()?;
let upvalue_count = self.read_u8()?;
let is_vararg = self.read_u8()? != 0;
let flags = self.read_u8()?;
let types_size = self.read_varint()? as usize;
let mut type_info = Vec::new();
let mut upvalue_types = Vec::new();
let mut local_types = Vec::new();
if types_size > 0 {
let type_info_size = self.read_varint()? as usize;
let typed_upvalue_count = self.read_varint()? as usize;
let typed_local_count = self.read_varint()? as usize;
type_info.extend_from_slice(self.read_bytes(type_info_size)?);
upvalue_types.reserve(typed_upvalue_count);
for _ in 0..typed_upvalue_count {
upvalue_types.push(self.read_u8()?);
}
local_types.reserve(typed_local_count);
for _ in 0..typed_local_count {
let ty = self.read_u8()?;
let register = self.read_u8()?;
let start_pc = self.read_varint()?;
let end_pc = start_pc + self.read_varint()?;
local_types.push(BytecodeTypedLocal {
ty,
register,
start_pc,
end_pc,
});
}
}
let code_word_count = self.read_varint()? as usize;
let code = self.read_code_words(code_word_count)?;
let constants = self.read_constants()?;
let proto_count = self.read_varint()? as usize;
let mut protos = Vec::with_capacity(proto_count);
for _ in 0..proto_count {
protos.push(self.read_varint()?);
}
let line_defined = self.read_varint()?;
let debug_name = self
.strings
.get_id(self.read_varint()?)?
.unwrap_or_default();
let lines = self.read_lines(code_word_count)?;
let (locals, upvalue_names) = self.read_debug_info()?;
if flags::LuauCallFeedback.get() {
let feedback_count = self.read_varint()?;
for _ in 0..feedback_count {
let _slot_type = self.read_u8()?;
let _pc = self.read_varint()?;
}
}
if flags::LuauCostModel.get() && flags & crate::opcodes::PROTO_FLAG_INLINABLE != 0 {
let _cost = self.read_varint64()?;
}
let mut function = BytecodeFunction {
max_stack_size,
num_params,
upvalue_count,
is_vararg,
flags,
type_info,
upvalue_types,
local_types,
blocks: Vec::new(),
instructions: Vec::new(),
constants,
immediates: Vec::new(),
phis: Vec::new(),
projections: Vec::new(),
registers: HashMap::new(),
table_shapes: self.table_shapes,
class_shapes: self.class_shapes,
entry_block: BytecodeBlockId::new(0),
exit_block: BytecodeBlockId::new(0),
pc_to_block: Vec::new(),
pc_to_instruction: Vec::new(),
protos,
line_defined,
debug_name,
lines,
locals,
upvalue_names,
};
build_function_graph(&mut function, &code)?;
Self::remap_local_pcs(&mut function, code_word_count as u32);
Ok(function)
}
fn remap_local_pcs(function: &mut BytecodeFunction<'_>, code_word_count: u32) {
let pc_to_graph_instruction = |pc: u32| {
function
.pc_to_instruction
.get(pc as usize)
.map(|instruction| instruction.index() as u32)
.unwrap_or(code_word_count)
};
for local in &mut function.local_types {
local.start_pc = pc_to_graph_instruction(local.start_pc);
local.end_pc = pc_to_graph_instruction(local.end_pc);
}
for local in &mut function.locals {
local.start_pc = pc_to_graph_instruction(local.start_pc);
local.end_pc = pc_to_graph_instruction(local.end_pc);
}
}
fn read_constants(
&mut self,
) -> Result<Vec<BytecodeFunctionConstant<'table>>, BytecodeReadError> {
let count = self.read_varint()? as usize;
let mut constants = Vec::with_capacity(count);
for _ in 0..count {
let offset = self.offset;
let tag = self.read_u8()?;
let constant = match tag {
tag if tag == BytecodeConstantTag::Nil as u8 => BytecodeFunctionConstant::Nil,
tag if tag == BytecodeConstantTag::Boolean as u8 => {
BytecodeFunctionConstant::Boolean(self.read_u8()? != 0)
}
tag if tag == BytecodeConstantTag::Number as u8 => {
BytecodeFunctionConstant::Number(self.read_f64()?)
}
tag if tag == BytecodeConstantTag::String as u8 => {
let id = self.read_varint()?;
let string = self
.strings
.get_id(id)?
.ok_or(BytecodeReadError::InvalidStringId { id })?;
BytecodeFunctionConstant::String(string)
}
tag if tag == BytecodeConstantTag::Import as u8 => {
BytecodeFunctionConstant::Import(self.read_u32()?)
}
tag if tag == BytecodeConstantTag::Table as u8 => {
let index = self.table_shapes.len() as u32;
let shape = self.read_table_shape(false)?;
self.table_shapes.push(shape);
BytecodeFunctionConstant::TableIndex(index)
}
tag if tag == BytecodeConstantTag::Closure as u8 => {
BytecodeFunctionConstant::Closure(self.read_varint()?)
}
tag if tag == BytecodeConstantTag::Vector as u8 => {
BytecodeFunctionConstant::Vector(BytecodeVector::new(
self.read_f32()?,
self.read_f32()?,
self.read_f32()?,
self.read_f32()?,
))
}
tag if tag == BytecodeConstantTag::VectorDouble as u8 => {
BytecodeFunctionConstant::VectorDouble(BytecodeVectorDouble::new(
self.read_f64()?,
self.read_f64()?,
self.read_f64()?,
self.read_f64()?,
))
}
tag if tag == BytecodeConstantTag::TableWithConstants as u8 => {
let index = self.table_shapes.len() as u32;
let shape = self.read_table_shape(true)?;
self.table_shapes.push(shape);
BytecodeFunctionConstant::TableIndex(index)
}
tag if tag == BytecodeConstantTag::Integer as u8 => {
BytecodeFunctionConstant::Integer(self.read_integer_constant()?)
}
tag if tag == BytecodeConstantTag::ClassShape as u8 => {
let index = self.class_shapes.len() as u32;
let class_name = self.read_varint()? as i32;
let property_count = self.read_varint()? as usize;
let method_count = self.read_varint()? as usize;
let mut property_names = Vec::with_capacity(property_count);
let mut method_names = Vec::with_capacity(method_count);
for _ in 0..property_count {
property_names.push(self.read_varint()? as i32);
}
for _ in 0..method_count {
method_names.push(self.read_varint()? as i32);
}
self.class_shapes.push(BytecodeClass {
class_name,
property_names,
method_names,
});
BytecodeFunctionConstant::ClassIndex(index)
}
tag => return Err(BytecodeReadError::InvalidConstantTag { tag, offset }),
};
constants.push(constant);
}
Ok(constants)
}
fn read_lines(&mut self, code_word_count: usize) -> Result<Vec<u32>, BytecodeReadError> {
let line_info = self.read_line_info(code_word_count)?;
if line_info.line_info.is_empty() {
return Ok(Vec::new());
}
Ok(line_info
.line_info
.into_iter()
.enumerate()
.map(|(pc, offset)| {
(line_info.abs_line_info[pc >> line_info.line_gap_log2] + i32::from(offset)) as u32
})
.collect())
}
fn read_debug_info(
&mut self,
) -> Result<(Vec<BytecodeDebugLocal<'table>>, Vec<&'table [u8]>), BytecodeReadError> {
if self.read_u8()? == 0 {
return Ok((Vec::new(), Vec::new()));
}
let local_count = self.read_varint()? as usize;
let mut locals = Vec::with_capacity(local_count);
for _ in 0..local_count {
let name_id = self.read_varint()?;
let name = self
.strings
.get_id(name_id)?
.ok_or(BytecodeReadError::InvalidStringId { id: name_id })?;
locals.push(BytecodeDebugLocal {
name,
start_pc: self.read_varint()?,
end_pc: self.read_varint()?,
register: self.read_u8()?,
});
}
let upvalue_count = self.read_varint()? as usize;
let mut upvalues = Vec::with_capacity(upvalue_count);
for _ in 0..upvalue_count {
let name_id = self.read_varint()?;
let name = self
.strings
.get_id(name_id)?
.ok_or(BytecodeReadError::InvalidStringId { id: name_id })?;
upvalues.push(name);
}
Ok((locals, upvalues))
}
fn read_u8(&mut self) -> Result<u8, BytecodeReadError> {
let offset = self.offset;
bytecode_wire::read_u8(self.data, &mut self.offset)
.ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<u8>()))
}
fn read_u32(&mut self) -> Result<u32, BytecodeReadError> {
let offset = self.offset;
bytecode_wire::read_u32(self.data, &mut self.offset)
.ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<u32>()))
}
fn read_i32(&mut self) -> Result<i32, BytecodeReadError> {
let offset = self.offset;
bytecode_wire::read_i32(self.data, &mut self.offset)
.ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<i32>()))
}
fn read_f32(&mut self) -> Result<f32, BytecodeReadError> {
let offset = self.offset;
bytecode_wire::read_f32(self.data, &mut self.offset)
.ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<f32>()))
}
fn read_f64(&mut self) -> Result<f64, BytecodeReadError> {
let offset = self.offset;
bytecode_wire::read_f64(self.data, &mut self.offset)
.ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<f64>()))
}
fn read_bytes(&mut self, len: usize) -> Result<&'data [u8], BytecodeReadError> {
let offset = self.offset;
bytecode_wire::read_bytes(self.data, &mut self.offset, len)
.ok_or_else(|| self.unexpected_eof(offset, len))
}
fn read_varint(&mut self) -> Result<u32, BytecodeReadError> {
let offset = self.offset;
bytecode_wire::read_varint(self.data, &mut self.offset)
.ok_or_else(|| self.unexpected_eof(offset, 1))
}
fn read_varint64(&mut self) -> Result<u64, BytecodeReadError> {
let offset = self.offset;
bytecode_wire::read_varint64(self.data, &mut self.offset)
.ok_or_else(|| self.unexpected_eof(offset, 1))
}
fn read_integer_constant(&mut self) -> Result<i64, BytecodeReadError> {
let negative = self.read_u8()? != 0;
let magnitude = self.read_varint64()?;
Ok(if negative {
(!magnitude).wrapping_add(1) as i64
} else {
magnitude as i64
})
}
fn read_code_words(
&mut self,
word_count: usize,
) -> Result<Vec<Instruction>, BytecodeReadError> {
let code_offset = self.offset;
let mut code = Vec::with_capacity(word_count);
for _ in 0..word_count {
code.push(Instruction::new(self.read_u32()?));
}
Self::validate_instruction_starts(&code, code_offset)?;
Ok(code)
}
fn validate_instruction_starts(
code: &[Instruction],
code_offset: usize,
) -> Result<(), BytecodeReadError> {
let mut pc = 0usize;
while pc < code.len() {
let offset = code_offset + pc * std::mem::size_of::<InstructionWord>();
let word = code[pc].word();
let opcode_byte = (word & 0xff) as u8;
let opcode =
Opcode::from_byte(opcode_byte).ok_or(BytecodeReadError::InvalidOpcode {
opcode: opcode_byte,
offset,
})?;
if pc + opcode.length() > code.len() {
return Err(BytecodeReadError::UnexpectedEof {
offset,
requested: opcode.length() * std::mem::size_of::<InstructionWord>(),
available: (code.len() - pc) * std::mem::size_of::<InstructionWord>(),
});
}
pc += opcode.length();
}
Ok(())
}
fn read_table_shape(&mut self, has_constants: bool) -> Result<TableShape, BytecodeReadError> {
let len = self.read_varint()? as usize;
let mut entries = Vec::with_capacity(len);
for _ in 0..len {
let key = self.read_varint()? as i32;
let value = if has_constants {
match self.read_i32()? {
-1 => None,
value => Some(value),
}
} else {
None
};
entries.push(TableShapeEntry { key, value });
}
Ok(TableShape::new(entries))
}
fn read_line_info(&mut self, code_len: usize) -> Result<BytecodeLineInfo, BytecodeReadError> {
if self.read_u8()? == 0 {
return Ok(BytecodeLineInfo::default());
}
let line_gap_log2 = self.read_u8()?;
let intervals = ((code_len.saturating_sub(1)) >> line_gap_log2) + 1;
let mut line_info = Vec::with_capacity(code_len);
let mut last_offset = 0u8;
for _ in 0..code_len {
last_offset = last_offset.wrapping_add(self.read_u8()?);
line_info.push(last_offset);
}
let mut abs_line_info = Vec::with_capacity(intervals);
let mut last_line = 0i32;
for _ in 0..intervals {
last_line = last_line.wrapping_add(self.read_i32()?);
abs_line_info.push(last_line);
}
Ok(BytecodeLineInfo {
line_info,
abs_line_info,
line_gap_log2,
})
}
fn unexpected_eof(&self, offset: usize, requested: usize) -> BytecodeReadError {
BytecodeReadError::UnexpectedEof {
offset,
requested,
available: self.data.len().saturating_sub(offset),
}
}
}
#[derive(Debug, Default)]
struct BytecodeLineInfo {
line_info: Vec<u8>,
abs_line_info: Vec<i32>,
line_gap_log2: u8,
}
#[derive(Debug, Clone, PartialEq)]
pub enum BytecodeFunctionConstant<'table> {
Nil,
Boolean(bool),
Number(f64),
Vector(BytecodeVector),
VectorDouble(BytecodeVectorDouble),
String(&'table [u8]),
Import(u32),
TableIndex(u32),
Closure(u32),
Integer(i64),
ClassIndex(u32),
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct BytecodeStringTable<'strings> {
strings: Vec<Cow<'strings, [u8]>>,
}
impl<'strings> BytecodeStringTable<'strings> {
pub fn new(strings: impl Into<Vec<Cow<'strings, [u8]>>>) -> Self {
Self {
strings: strings.into(),
}
}
pub fn len(&self) -> usize {
self.strings.len()
}
pub fn is_empty(&self) -> bool {
self.strings.is_empty()
}
pub fn get(&self, index: usize) -> Option<&[u8]> {
self.strings.get(index).map(Cow::as_ref)
}
pub fn iter(&self) -> impl Iterator<Item = &[u8]> + '_ {
self.strings.iter().map(Cow::as_ref)
}
pub fn get_id(&self, id: u32) -> Result<Option<&[u8]>, BytecodeReadError> {
if id == 0 {
return Ok(None);
}
self.get(id as usize - 1)
.map(Some)
.ok_or(BytecodeReadError::InvalidStringId { id })
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BytecodeDebugLocal<'table> {
pub name: &'table [u8],
pub register: u8,
pub start_pc: u32,
pub end_pc: u32,
}