1use crate::builder::BytecodeBuilder;
2use crate::error::BytecodeReadError;
3use crate::graph::{
4 BytecodeBlock, BytecodeBlockId, BytecodeImmediate, BytecodeImmediateId, BytecodeInstruction,
5 BytecodeInstructionId, BytecodeOperand, BytecodePhi, BytecodePhiId, BytecodeProjection,
6 BytecodeProjectionId, BytecodeWriteError, InstructionPc, build_function_graph,
7 encode_function_bytecode,
8};
9use crate::model::{
10 BytecodeClass, BytecodeTypedLocal, BytecodeVector, BytecodeVectorDouble, Instruction,
11 InstructionWord, Register, TableShape, TableShapeEntry,
12};
13use crate::opcodes::{BytecodeConstantTag, Opcode};
14use luau_common::{bytecode_wire, flags};
15use std::borrow::Cow;
16use std::collections::HashMap;
17
18#[derive(Debug, Clone, PartialEq)]
19pub struct BytecodeFunction<'table> {
20 pub max_stack_size: u8,
21 pub num_params: u8,
22 pub upvalue_count: u8,
23 pub is_vararg: bool,
24 pub flags: u8,
25 pub type_info: Vec<u8>,
26 pub upvalue_types: Vec<u8>,
27 pub local_types: Vec<BytecodeTypedLocal>,
28 pub blocks: Vec<BytecodeBlock>,
29 pub instructions: Vec<BytecodeInstruction>,
30 pub constants: Vec<BytecodeFunctionConstant<'table>>,
31 pub immediates: Vec<BytecodeImmediate>,
32 pub phis: Vec<BytecodePhi>,
33 pub projections: Vec<BytecodeProjection>,
34 pub registers: HashMap<BytecodeOperand, Register>,
35 pub table_shapes: Vec<TableShape>,
36 pub class_shapes: Vec<BytecodeClass>,
37 pub entry_block: BytecodeBlockId,
38 pub exit_block: BytecodeBlockId,
39 pub pc_to_block: Vec<BytecodeBlockId>,
40 pub pc_to_instruction: Vec<BytecodeInstructionId>,
41 pub protos: Vec<u32>,
42 pub line_defined: u32,
43 pub debug_name: &'table [u8],
44 pub lines: Vec<u32>,
45 pub locals: Vec<BytecodeDebugLocal<'table>>,
46 pub upvalue_names: Vec<&'table [u8]>,
47}
48
49impl<'table> BytecodeFunction<'table> {
50 pub fn from_function_bytecode<'strings>(
51 data: &[u8],
52 strings: &'table BytecodeStringTable<'strings>,
53 ) -> Result<Self, BytecodeReadError> {
54 BytecodeFunctionReader::new(data, strings).read()
55 }
56
57 pub fn to_function_bytecode(&mut self) -> Result<Vec<u8>, BytecodeWriteError> {
62 let mut builder = BytecodeBuilder::new();
63 encode_function_bytecode(&mut builder, self)
64 }
65
66 pub fn entry(&self) -> BytecodeBlockId {
67 self.entry_block
68 }
69
70 pub fn exit(&self) -> BytecodeBlockId {
71 self.exit_block
72 }
73
74 pub fn blocks(&self) -> &[BytecodeBlock] {
75 &self.blocks
76 }
77
78 pub fn block(&self, id: BytecodeBlockId) -> &BytecodeBlock {
79 &self.blocks[id.index()]
80 }
81
82 pub fn graph_instruction(&self, id: BytecodeInstructionId) -> &BytecodeInstruction {
83 &self.instructions[id.index()]
84 }
85
86 pub fn block_for_pc(&self, pc: InstructionPc) -> Option<BytecodeBlockId> {
87 self.pc_to_block.get(pc.index()).copied()
88 }
89
90 pub fn instruction_for_pc(&self, pc: InstructionPc) -> Option<BytecodeInstructionId> {
91 self.pc_to_instruction.get(pc.index()).copied()
92 }
93
94 pub fn immediate(&self, id: BytecodeImmediateId) -> &BytecodeImmediate {
95 &self.immediates[id.index()]
96 }
97
98 pub fn phi(&self, id: BytecodePhiId) -> &BytecodePhi {
99 &self.phis[id.index()]
100 }
101
102 pub fn projection(&self, id: BytecodeProjectionId) -> &BytecodeProjection {
103 &self.projections[id.index()]
104 }
105}
106
107struct BytecodeFunctionReader<'data, 'table, 'strings> {
108 data: &'data [u8],
109 offset: usize,
110 strings: &'table BytecodeStringTable<'strings>,
111 table_shapes: Vec<TableShape>,
112 class_shapes: Vec<BytecodeClass>,
113}
114
115impl<'data, 'table, 'strings> BytecodeFunctionReader<'data, 'table, 'strings> {
116 fn new(data: &'data [u8], strings: &'table BytecodeStringTable<'strings>) -> Self {
117 Self {
118 data,
119 offset: 0,
120 strings,
121 table_shapes: Vec::new(),
122 class_shapes: Vec::new(),
123 }
124 }
125
126 fn read(mut self) -> Result<BytecodeFunction<'table>, BytecodeReadError> {
127 let max_stack_size = self.read_u8()?;
128 let num_params = self.read_u8()?;
129 let upvalue_count = self.read_u8()?;
130 let is_vararg = self.read_u8()? != 0;
131 let flags = self.read_u8()?;
132
133 let types_size = self.read_varint()? as usize;
134 let mut type_info = Vec::new();
135 let mut upvalue_types = Vec::new();
136 let mut local_types = Vec::new();
137
138 if types_size > 0 {
139 let type_info_size = self.read_varint()? as usize;
140 let typed_upvalue_count = self.read_varint()? as usize;
141 let typed_local_count = self.read_varint()? as usize;
142
143 type_info.extend_from_slice(self.read_bytes(type_info_size)?);
144
145 upvalue_types.reserve(typed_upvalue_count);
146 for _ in 0..typed_upvalue_count {
147 upvalue_types.push(self.read_u8()?);
148 }
149
150 local_types.reserve(typed_local_count);
151 for _ in 0..typed_local_count {
152 let ty = self.read_u8()?;
153 let register = self.read_u8()?;
154 let start_pc = self.read_varint()?;
155 let end_pc = start_pc + self.read_varint()?;
156 local_types.push(BytecodeTypedLocal {
157 ty,
158 register,
159 start_pc,
160 end_pc,
161 });
162 }
163 }
164
165 let code_word_count = self.read_varint()? as usize;
166 let code = self.read_code_words(code_word_count)?;
167 let constants = self.read_constants()?;
168
169 let proto_count = self.read_varint()? as usize;
170 let mut protos = Vec::with_capacity(proto_count);
171 for _ in 0..proto_count {
172 protos.push(self.read_varint()?);
173 }
174
175 let line_defined = self.read_varint()?;
176 let debug_name = self
177 .strings
178 .get_id(self.read_varint()?)?
179 .unwrap_or_default();
180 let lines = self.read_lines(code_word_count)?;
181 let (locals, upvalue_names) = self.read_debug_info()?;
182
183 if flags::LuauCallFeedback.get() {
184 let feedback_count = self.read_varint()?;
185 for _ in 0..feedback_count {
186 let _slot_type = self.read_u8()?;
187 let _pc = self.read_varint()?;
188 }
189 }
190
191 if flags::LuauCostModel.get() && flags & crate::opcodes::PROTO_FLAG_INLINABLE != 0 {
192 let _cost = self.read_varint64()?;
193 }
194
195 let mut function = BytecodeFunction {
196 max_stack_size,
197 num_params,
198 upvalue_count,
199 is_vararg,
200 flags,
201 type_info,
202 upvalue_types,
203 local_types,
204 blocks: Vec::new(),
205 instructions: Vec::new(),
206 constants,
207 immediates: Vec::new(),
208 phis: Vec::new(),
209 projections: Vec::new(),
210 registers: HashMap::new(),
211 table_shapes: self.table_shapes,
212 class_shapes: self.class_shapes,
213 entry_block: BytecodeBlockId::new(0),
214 exit_block: BytecodeBlockId::new(0),
215 pc_to_block: Vec::new(),
216 pc_to_instruction: Vec::new(),
217 protos,
218 line_defined,
219 debug_name,
220 lines,
221 locals,
222 upvalue_names,
223 };
224
225 build_function_graph(&mut function, &code)?;
226 Self::remap_local_pcs(&mut function, code_word_count as u32);
227 Ok(function)
228 }
229
230 fn remap_local_pcs(function: &mut BytecodeFunction<'_>, code_word_count: u32) {
231 let pc_to_graph_instruction = |pc: u32| {
232 function
233 .pc_to_instruction
234 .get(pc as usize)
235 .map(|instruction| instruction.index() as u32)
236 .unwrap_or(code_word_count)
237 };
238
239 for local in &mut function.local_types {
240 local.start_pc = pc_to_graph_instruction(local.start_pc);
241 local.end_pc = pc_to_graph_instruction(local.end_pc);
242 }
243
244 for local in &mut function.locals {
245 local.start_pc = pc_to_graph_instruction(local.start_pc);
246 local.end_pc = pc_to_graph_instruction(local.end_pc);
247 }
248 }
249
250 fn read_constants(
251 &mut self,
252 ) -> Result<Vec<BytecodeFunctionConstant<'table>>, BytecodeReadError> {
253 let count = self.read_varint()? as usize;
254 let mut constants = Vec::with_capacity(count);
255
256 for _ in 0..count {
257 let offset = self.offset;
258 let tag = self.read_u8()?;
259 let constant = match tag {
260 tag if tag == BytecodeConstantTag::Nil as u8 => BytecodeFunctionConstant::Nil,
261 tag if tag == BytecodeConstantTag::Boolean as u8 => {
262 BytecodeFunctionConstant::Boolean(self.read_u8()? != 0)
263 }
264 tag if tag == BytecodeConstantTag::Number as u8 => {
265 BytecodeFunctionConstant::Number(self.read_f64()?)
266 }
267 tag if tag == BytecodeConstantTag::String as u8 => {
268 let id = self.read_varint()?;
269 let string = self
270 .strings
271 .get_id(id)?
272 .ok_or(BytecodeReadError::InvalidStringId { id })?;
273 BytecodeFunctionConstant::String(string)
274 }
275 tag if tag == BytecodeConstantTag::Import as u8 => {
276 BytecodeFunctionConstant::Import(self.read_u32()?)
277 }
278 tag if tag == BytecodeConstantTag::Table as u8 => {
279 let index = self.table_shapes.len() as u32;
280 let shape = self.read_table_shape(false)?;
281 self.table_shapes.push(shape);
282 BytecodeFunctionConstant::TableIndex(index)
283 }
284 tag if tag == BytecodeConstantTag::Closure as u8 => {
285 BytecodeFunctionConstant::Closure(self.read_varint()?)
286 }
287 tag if tag == BytecodeConstantTag::Vector as u8 => {
288 BytecodeFunctionConstant::Vector(BytecodeVector::new(
289 self.read_f32()?,
290 self.read_f32()?,
291 self.read_f32()?,
292 self.read_f32()?,
293 ))
294 }
295 tag if tag == BytecodeConstantTag::VectorDouble as u8 => {
296 BytecodeFunctionConstant::VectorDouble(BytecodeVectorDouble::new(
297 self.read_f64()?,
298 self.read_f64()?,
299 self.read_f64()?,
300 self.read_f64()?,
301 ))
302 }
303 tag if tag == BytecodeConstantTag::TableWithConstants as u8 => {
304 let index = self.table_shapes.len() as u32;
305 let shape = self.read_table_shape(true)?;
306 self.table_shapes.push(shape);
307 BytecodeFunctionConstant::TableIndex(index)
308 }
309 tag if tag == BytecodeConstantTag::Integer as u8 => {
310 BytecodeFunctionConstant::Integer(self.read_integer_constant()?)
311 }
312 tag if tag == BytecodeConstantTag::ClassShape as u8 => {
313 let index = self.class_shapes.len() as u32;
314 let class_name = self.read_varint()? as i32;
315 let property_count = self.read_varint()? as usize;
316 let method_count = self.read_varint()? as usize;
317 let mut property_names = Vec::with_capacity(property_count);
318 let mut method_names = Vec::with_capacity(method_count);
319 for _ in 0..property_count {
320 property_names.push(self.read_varint()? as i32);
321 }
322 for _ in 0..method_count {
323 method_names.push(self.read_varint()? as i32);
324 }
325 self.class_shapes.push(BytecodeClass {
326 class_name,
327 property_names,
328 method_names,
329 });
330 BytecodeFunctionConstant::ClassIndex(index)
331 }
332 tag => return Err(BytecodeReadError::InvalidConstantTag { tag, offset }),
333 };
334 constants.push(constant);
335 }
336
337 Ok(constants)
338 }
339
340 fn read_lines(&mut self, code_word_count: usize) -> Result<Vec<u32>, BytecodeReadError> {
341 let line_info = self.read_line_info(code_word_count)?;
342 if line_info.line_info.is_empty() {
343 return Ok(Vec::new());
344 }
345
346 Ok(line_info
347 .line_info
348 .into_iter()
349 .enumerate()
350 .map(|(pc, offset)| {
351 (line_info.abs_line_info[pc >> line_info.line_gap_log2] + i32::from(offset)) as u32
352 })
353 .collect())
354 }
355
356 fn read_debug_info(
357 &mut self,
358 ) -> Result<(Vec<BytecodeDebugLocal<'table>>, Vec<&'table [u8]>), BytecodeReadError> {
359 if self.read_u8()? == 0 {
360 return Ok((Vec::new(), Vec::new()));
361 }
362
363 let local_count = self.read_varint()? as usize;
364 let mut locals = Vec::with_capacity(local_count);
365 for _ in 0..local_count {
366 let name_id = self.read_varint()?;
367 let name = self
368 .strings
369 .get_id(name_id)?
370 .ok_or(BytecodeReadError::InvalidStringId { id: name_id })?;
371 locals.push(BytecodeDebugLocal {
372 name,
373 start_pc: self.read_varint()?,
374 end_pc: self.read_varint()?,
375 register: self.read_u8()?,
376 });
377 }
378
379 let upvalue_count = self.read_varint()? as usize;
380 let mut upvalues = Vec::with_capacity(upvalue_count);
381 for _ in 0..upvalue_count {
382 let name_id = self.read_varint()?;
383 let name = self
384 .strings
385 .get_id(name_id)?
386 .ok_or(BytecodeReadError::InvalidStringId { id: name_id })?;
387 upvalues.push(name);
388 }
389
390 Ok((locals, upvalues))
391 }
392
393 fn read_u8(&mut self) -> Result<u8, BytecodeReadError> {
394 let offset = self.offset;
395 bytecode_wire::read_u8(self.data, &mut self.offset)
396 .ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<u8>()))
397 }
398
399 fn read_u32(&mut self) -> Result<u32, BytecodeReadError> {
400 let offset = self.offset;
401 bytecode_wire::read_u32(self.data, &mut self.offset)
402 .ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<u32>()))
403 }
404
405 fn read_i32(&mut self) -> Result<i32, BytecodeReadError> {
406 let offset = self.offset;
407 bytecode_wire::read_i32(self.data, &mut self.offset)
408 .ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<i32>()))
409 }
410
411 fn read_f32(&mut self) -> Result<f32, BytecodeReadError> {
412 let offset = self.offset;
413 bytecode_wire::read_f32(self.data, &mut self.offset)
414 .ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<f32>()))
415 }
416
417 fn read_f64(&mut self) -> Result<f64, BytecodeReadError> {
418 let offset = self.offset;
419 bytecode_wire::read_f64(self.data, &mut self.offset)
420 .ok_or_else(|| self.unexpected_eof(offset, std::mem::size_of::<f64>()))
421 }
422
423 fn read_bytes(&mut self, len: usize) -> Result<&'data [u8], BytecodeReadError> {
424 let offset = self.offset;
425 bytecode_wire::read_bytes(self.data, &mut self.offset, len)
426 .ok_or_else(|| self.unexpected_eof(offset, len))
427 }
428
429 fn read_varint(&mut self) -> Result<u32, BytecodeReadError> {
430 let offset = self.offset;
431 bytecode_wire::read_varint(self.data, &mut self.offset)
432 .ok_or_else(|| self.unexpected_eof(offset, 1))
433 }
434
435 fn read_varint64(&mut self) -> Result<u64, BytecodeReadError> {
436 let offset = self.offset;
437 bytecode_wire::read_varint64(self.data, &mut self.offset)
438 .ok_or_else(|| self.unexpected_eof(offset, 1))
439 }
440
441 fn read_integer_constant(&mut self) -> Result<i64, BytecodeReadError> {
442 let negative = self.read_u8()? != 0;
443 let magnitude = self.read_varint64()?;
444
445 Ok(if negative {
446 (!magnitude).wrapping_add(1) as i64
447 } else {
448 magnitude as i64
449 })
450 }
451
452 fn read_code_words(
453 &mut self,
454 word_count: usize,
455 ) -> Result<Vec<Instruction>, BytecodeReadError> {
456 let code_offset = self.offset;
457 let mut code = Vec::with_capacity(word_count);
458
459 for _ in 0..word_count {
460 code.push(Instruction::new(self.read_u32()?));
461 }
462
463 Self::validate_instruction_starts(&code, code_offset)?;
464 Ok(code)
465 }
466
467 fn validate_instruction_starts(
468 code: &[Instruction],
469 code_offset: usize,
470 ) -> Result<(), BytecodeReadError> {
471 let mut pc = 0usize;
472 while pc < code.len() {
473 let offset = code_offset + pc * std::mem::size_of::<InstructionWord>();
474 let word = code[pc].word();
475 let opcode_byte = (word & 0xff) as u8;
476 let opcode =
477 Opcode::from_byte(opcode_byte).ok_or(BytecodeReadError::InvalidOpcode {
478 opcode: opcode_byte,
479 offset,
480 })?;
481
482 if pc + opcode.length() > code.len() {
483 return Err(BytecodeReadError::UnexpectedEof {
484 offset,
485 requested: opcode.length() * std::mem::size_of::<InstructionWord>(),
486 available: (code.len() - pc) * std::mem::size_of::<InstructionWord>(),
487 });
488 }
489
490 pc += opcode.length();
491 }
492
493 Ok(())
494 }
495
496 fn read_table_shape(&mut self, has_constants: bool) -> Result<TableShape, BytecodeReadError> {
497 let len = self.read_varint()? as usize;
498 let mut entries = Vec::with_capacity(len);
499
500 for _ in 0..len {
501 let key = self.read_varint()? as i32;
502 let value = if has_constants {
503 match self.read_i32()? {
504 -1 => None,
505 value => Some(value),
506 }
507 } else {
508 None
509 };
510 entries.push(TableShapeEntry { key, value });
511 }
512
513 Ok(TableShape::new(entries))
514 }
515
516 fn read_line_info(&mut self, code_len: usize) -> Result<BytecodeLineInfo, BytecodeReadError> {
517 if self.read_u8()? == 0 {
518 return Ok(BytecodeLineInfo::default());
519 }
520
521 let line_gap_log2 = self.read_u8()?;
522 let intervals = ((code_len.saturating_sub(1)) >> line_gap_log2) + 1;
523 let mut line_info = Vec::with_capacity(code_len);
524 let mut last_offset = 0u8;
525
526 for _ in 0..code_len {
527 last_offset = last_offset.wrapping_add(self.read_u8()?);
528 line_info.push(last_offset);
529 }
530
531 let mut abs_line_info = Vec::with_capacity(intervals);
532 let mut last_line = 0i32;
533 for _ in 0..intervals {
534 last_line = last_line.wrapping_add(self.read_i32()?);
535 abs_line_info.push(last_line);
536 }
537
538 Ok(BytecodeLineInfo {
539 line_info,
540 abs_line_info,
541 line_gap_log2,
542 })
543 }
544
545 fn unexpected_eof(&self, offset: usize, requested: usize) -> BytecodeReadError {
546 BytecodeReadError::UnexpectedEof {
547 offset,
548 requested,
549 available: self.data.len().saturating_sub(offset),
550 }
551 }
552}
553
554#[derive(Debug, Default)]
555struct BytecodeLineInfo {
556 line_info: Vec<u8>,
557 abs_line_info: Vec<i32>,
558 line_gap_log2: u8,
559}
560
561#[derive(Debug, Clone, PartialEq)]
562pub enum BytecodeFunctionConstant<'table> {
563 Nil,
564 Boolean(bool),
565 Number(f64),
566 Vector(BytecodeVector),
567 VectorDouble(BytecodeVectorDouble),
568 String(&'table [u8]),
569 Import(u32),
570 TableIndex(u32),
571 Closure(u32),
572 Integer(i64),
573 ClassIndex(u32),
574}
575
576#[derive(Debug, Clone, Default, PartialEq, Eq)]
577pub struct BytecodeStringTable<'strings> {
578 strings: Vec<Cow<'strings, [u8]>>,
579}
580
581impl<'strings> BytecodeStringTable<'strings> {
582 pub fn new(strings: impl Into<Vec<Cow<'strings, [u8]>>>) -> Self {
583 Self {
584 strings: strings.into(),
585 }
586 }
587
588 pub fn len(&self) -> usize {
589 self.strings.len()
590 }
591
592 pub fn is_empty(&self) -> bool {
593 self.strings.is_empty()
594 }
595
596 pub fn get(&self, index: usize) -> Option<&[u8]> {
597 self.strings.get(index).map(Cow::as_ref)
598 }
599
600 pub fn iter(&self) -> impl Iterator<Item = &[u8]> + '_ {
601 self.strings.iter().map(Cow::as_ref)
602 }
603
604 pub fn get_id(&self, id: u32) -> Result<Option<&[u8]>, BytecodeReadError> {
605 if id == 0 {
606 return Ok(None);
607 }
608
609 self.get(id as usize - 1)
610 .map(Some)
611 .ok_or(BytecodeReadError::InvalidStringId { id })
612 }
613}
614
615#[derive(Debug, Clone, Copy, PartialEq, Eq)]
616pub struct BytecodeDebugLocal<'table> {
617 pub name: &'table [u8],
618 pub register: u8,
619 pub start_pc: u32,
620 pub end_pc: u32,
621}