Skip to main content

luau_bytecode/builder/
mod.rs

1mod debug;
2mod emit;
3mod serialize;
4mod support;
5
6use crate::model::{BytecodeClass, BytecodeUserdataType, InstructionWord};
7use luau_common::{BString, DenseHashMap};
8
9use support::{
10    BytecodeBuilderFunction, BytecodeBuilderScratch, BytecodeStringHasher, StoredDebugRemark,
11};
12pub use support::{BytecodeDumpFlags, BytecodeEncoder};
13
14const MAX_CONSTANT_COUNT: usize = 1 << 23;
15
16pub struct BytecodeBuilder<'src> {
17    functions: Vec<BytecodeBuilderFunction>,
18    scratch: BytecodeBuilderScratch<'src>,
19    class_shapes: Vec<BytecodeClass>,
20    userdata_types: Vec<BytecodeUserdataType>,
21    string_index: DenseHashMap<BytecodeStringRef<'src>, u32, BytecodeStringHasher>,
22    debug_strings: Vec<BytecodeStringRef<'src>>,
23    bytecode: Vec<u8>,
24    encoder: Option<Box<dyn BytecodeEncoder>>,
25    current_function: Option<usize>,
26    main: Option<usize>,
27    total_instruction_count: usize,
28    current_line: i32,
29    dump_flags: BytecodeDumpFlags,
30    dump_enabled: bool,
31    dump_source: Vec<BString>,
32    dump_remarks: Vec<(i32, BString)>,
33}
34
35#[derive(Debug, Clone, Copy)]
36pub struct BytecodeStringRef<'src> {
37    data: *const u8,
38    length: usize,
39    _marker: std::marker::PhantomData<&'src [u8]>,
40}
41
42impl<'src> BytecodeStringRef<'src> {
43    pub const fn empty() -> Self {
44        Self {
45            data: std::ptr::null(),
46            length: 0,
47            _marker: std::marker::PhantomData,
48        }
49    }
50
51    pub fn as_bytes(&self) -> &'src [u8] {
52        if self.data.is_null() {
53            &[]
54        } else {
55            unsafe { std::slice::from_raw_parts(self.data, self.length) }
56        }
57    }
58}
59
60impl PartialEq for BytecodeStringRef<'_> {
61    fn eq(&self, other: &Self) -> bool {
62        if !self.data.is_null() && !other.data.is_null() {
63            self.length == other.length && self.as_bytes() == other.as_bytes()
64        } else {
65            self.data == other.data
66        }
67    }
68}
69
70impl Eq for BytecodeStringRef<'_> {}
71
72impl<'src> From<&'src [u8]> for BytecodeStringRef<'src> {
73    fn from(value: &'src [u8]) -> Self {
74        Self {
75            data: value.as_ptr(),
76            length: value.len(),
77            _marker: std::marker::PhantomData,
78        }
79    }
80}
81
82impl<'src, const N: usize> From<&'src [u8; N]> for BytecodeStringRef<'src> {
83    fn from(value: &'src [u8; N]) -> Self {
84        Self::from(value.as_slice())
85    }
86}
87
88impl<'src> Default for BytecodeBuilder<'src> {
89    fn default() -> Self {
90        Self {
91            functions: Vec::new(),
92            scratch: BytecodeBuilderScratch::default(),
93            class_shapes: Vec::new(),
94            userdata_types: Vec::new(),
95            string_index: DenseHashMap::new(BytecodeStringRef::empty()),
96            debug_strings: Vec::new(),
97            bytecode: Vec::new(),
98            encoder: None,
99            current_function: None,
100            main: None,
101            total_instruction_count: 0,
102            current_line: 0,
103            dump_flags: BytecodeDumpFlags::default(),
104            dump_enabled: false,
105            dump_source: Vec::new(),
106            dump_remarks: Vec::new(),
107        }
108    }
109}
110
111impl<'src> BytecodeBuilder<'src> {
112    pub fn new() -> Self {
113        Self::default()
114    }
115
116    pub fn with_encoder(encoder: impl BytecodeEncoder + 'static) -> Self {
117        Self {
118            encoder: Some(Box::new(encoder)),
119            ..Self::default()
120        }
121    }
122
123    pub fn begin_function(&mut self, num_params: u8, is_vararg: bool) -> usize {
124        debug_assert!(
125            self.current_function.is_none(),
126            "BytecodeBuilder::begin_function requires no active function"
127        );
128
129        let id = self.functions.len();
130        self.functions
131            .push(BytecodeBuilderFunction::new(num_params, is_vararg));
132        self.current_function = Some(id);
133        self.current_line = 0;
134        id
135    }
136
137    pub fn function_count(&self) -> usize {
138        self.functions.len()
139    }
140
141    pub fn clear_string_table(&mut self) {
142        self.string_index.clear();
143    }
144
145    pub fn end_function(
146        &mut self,
147        max_stack_size: u8,
148        upvalue_count: u8,
149        flags: u8,
150        cost: u64,
151    ) -> usize {
152        let id = self
153            .current_function
154            .take()
155            .expect("BytecodeBuilder::end_function requires an active function");
156        self.functions[id].max_stack_size = max_stack_size;
157        self.functions[id].upvalue_count = upvalue_count;
158        self.functions[id].flags |= flags;
159        self.functions[id].cost = cost;
160        #[cfg(debug_assertions)]
161        self.scratch.validate(&self.functions[id], &self.functions);
162        let (dump, dump_instruction_offsets) = self.dump_current_function(id);
163        self.functions[id].dump = dump;
164        self.functions[id].dump_instruction_offsets = dump_instruction_offsets;
165        if let Some(encoder) = self.encoder.as_deref() {
166            let code = &mut self.scratch.code;
167            // Safety: `Instruction` is `repr(transparent)` over `InstructionWord`, so the live
168            // instruction buffer has the exact layout the encoder expects.
169            let words = unsafe {
170                std::slice::from_raw_parts_mut(
171                    code.as_mut_ptr().cast::<InstructionWord>(),
172                    code.len(),
173                )
174            };
175            encoder.encode(words);
176        }
177        self.functions[id].data = self.function_data(id);
178        self.total_instruction_count += self.scratch.code.len();
179        self.scratch.clear();
180        id
181    }
182
183    pub fn add_child_function(&mut self, id: u32) -> Option<i16> {
184        self.current_function
185            .expect("BytecodeBuilder::add_child_function requires an active function");
186
187        if let Some(index) = self.scratch.child_function_map.get(&id) {
188            return Some(*index);
189        }
190
191        let index = i16::try_from(self.scratch.child_functions.len()).ok()?;
192        self.scratch.child_functions.push(id);
193        let (_, fresh) = self.scratch.child_function_map.insert(id, index);
194        debug_assert!(fresh);
195        Some(index)
196    }
197
198    pub fn set_main_function(&mut self, id: usize) {
199        debug_assert!(
200            id < self.functions.len(),
201            "BytecodeBuilder::set_main_function requires a valid function id"
202        );
203        self.main = Some(id);
204    }
205}