lust/bytecode/compiler/
module.rs1use super::*;
2impl Compiler {
3 pub fn compile_module(&mut self, items: &[Item]) -> Result<Vec<Function>> {
4 let mut script_stmts: Vec<Stmt> = Vec::new();
5 for item in items {
6 match &item.kind {
7 ItemKind::Module { name, items } => {
8 let prev = self.current_module.clone();
9 self.current_module = Some(name.clone());
10 for ni in items {
11 match &ni.kind {
12 ItemKind::Script(stmts) => {
13 script_stmts.extend_from_slice(stmts);
14 {
15 let locals_entry =
16 self.module_locals.entry(name.clone()).or_default();
17 for local in Self::collect_top_level_locals(stmts) {
18 locals_entry.insert(local);
19 }
20 }
21 }
22
23 ItemKind::Function(func_def) => {
24 if Self::is_module_init_function(func_def, name) {
25 let locals_entry =
26 self.module_locals.entry(name.clone()).or_default();
27 for local in Self::collect_top_level_locals(&func_def.body) {
28 locals_entry.insert(local);
29 }
30 }
31
32 let func_idx = self.functions.len();
33 let function = Function::new(
34 &func_def.name,
35 func_def.params.len() as u8,
36 func_def.is_method,
37 );
38 self.function_table.insert(func_def.name.clone(), func_idx);
39 self.functions.push(function);
40 }
41
42 ItemKind::Struct(_) | ItemKind::Enum(_) => {}
43 ItemKind::Trait(trait_def) => {
44 self.trait_names.insert(trait_def.name.clone());
45 }
46
47 ItemKind::Impl(impl_block) => {
48 let type_name = match &impl_block.target_type.kind {
49 crate::ast::TypeKind::Named(name) => {
50 self.resolve_type_name(name)
51 }
52
53 _ => {
54 return Err(LustError::CompileError(
55 "Impl block target must be a named type".to_string(),
56 ))
57 }
58 };
59 if let Some(trait_name) = &impl_block.trait_name {
60 let resolved_trait = self.resolve_type_name(trait_name);
61 self.trait_impls.push((type_name.clone(), resolved_trait));
62 }
63
64 for method in &impl_block.methods {
65 let func_idx = self.functions.len();
66 let has_self =
67 method.params.iter().any(|p| p.is_self || p.name == "self");
68 let mangled_name =
69 if method.name.contains(':') || method.name.contains('.') {
70 method.name.clone()
71 } else if has_self {
72 format!("{}:{}", type_name, method.name)
73 } else {
74 format!("{}.{}", type_name, method.name)
75 };
76 let function = Function::new(
77 &mangled_name,
78 method.params.len() as u8,
79 has_self,
80 );
81 self.function_table.insert(mangled_name, func_idx);
82 self.functions.push(function);
83 }
84 }
85
86 ItemKind::Extern { items, .. } => {
87 for ext in items {
88 match ext {
89 ExternItem::Function { name, .. } => {
90 self.record_extern_function(name);
91 }
92 }
93 }
94 }
95
96 ItemKind::TypeAlias { .. }
97 | ItemKind::Use { .. }
98 | ItemKind::Const { .. }
99 | ItemKind::Static { .. } => {}
100 ItemKind::Module { .. } => {}
101 }
102 }
103
104 self.current_module = prev;
105 }
106
107 ItemKind::Script(stmts) => {
108 script_stmts.extend_from_slice(stmts);
109 let module_key = self
110 .entry_module
111 .clone()
112 .unwrap_or_else(|| "__root".to_string());
113 let locals_entry = self.module_locals.entry(module_key).or_default();
114 for local in Self::collect_top_level_locals(stmts) {
115 locals_entry.insert(local);
116 }
117 }
118
119 ItemKind::Function(func_def) => {
120 let func_idx = self.functions.len();
121 let function = Function::new(
122 &func_def.name,
123 func_def.params.len() as u8,
124 func_def.is_method,
125 );
126 self.function_table.insert(func_def.name.clone(), func_idx);
127 self.functions.push(function);
128 }
129
130 ItemKind::Struct(_) | ItemKind::Enum(_) => {}
131 ItemKind::Trait(trait_def) => {
132 self.trait_names.insert(trait_def.name.clone());
133 }
134
135 ItemKind::Impl(impl_block) => {
136 let type_name = match &impl_block.target_type.kind {
137 crate::ast::TypeKind::Named(name) => self.resolve_type_name(name),
138 _ => {
139 return Err(LustError::CompileError(
140 "Impl block target must be a named type".to_string(),
141 ));
142 }
143 };
144 if let Some(trait_name) = &impl_block.trait_name {
145 let resolved_trait = self.resolve_type_name(trait_name);
146 self.trait_impls.push((type_name.clone(), resolved_trait));
147 }
148
149 for method in &impl_block.methods {
150 let func_idx = self.functions.len();
151 let has_self = method.params.iter().any(|p| p.is_self || p.name == "self");
152 let mangled_name = if method.name.contains(':') || method.name.contains('.')
153 {
154 method.name.clone()
155 } else if has_self {
156 format!("{}:{}", type_name, method.name)
157 } else {
158 format!("{}.{}", type_name, method.name)
159 };
160 let function =
161 Function::new(&mangled_name, method.params.len() as u8, has_self);
162 self.function_table.insert(mangled_name, func_idx);
163 self.functions.push(function);
164 }
165 }
166
167 ItemKind::Extern { items, .. } => {
168 for ext in items {
169 match ext {
170 ExternItem::Function { name, .. } => {
171 self.record_extern_function(name);
172 }
173 }
174 }
175 }
176
177 ItemKind::TypeAlias { .. }
178 | ItemKind::Use { .. }
179 | ItemKind::Const { .. }
180 | ItemKind::Static { .. } => {}
181 }
182 }
183
184 let script_func_idx = self.functions.len();
185 let script_func = Function::new("__script", 0, false);
186 self.function_table
187 .insert("__script".to_string(), script_func_idx);
188 self.functions.push(script_func);
189 let mut func_idx = 0;
190 for item in items {
191 match &item.kind {
192 ItemKind::Module { name, items } => {
193 let prev = self.current_module.clone();
194 self.current_module = Some(name.clone());
195 for ni in items {
196 match &ni.kind {
197 ItemKind::Function(func_def) => {
198 self.compile_function(func_idx, func_def)?;
199 func_idx += 1;
200 }
201
202 ItemKind::Impl(impl_block) => {
203 for method in &impl_block.methods {
204 self.compile_function(func_idx, method)?;
205 func_idx += 1;
206 }
207 }
208
209 ItemKind::Module { .. } => {}
210 _ => {}
211 }
212 }
213
214 self.current_module = prev;
215 }
216
217 ItemKind::Function(func_def) => {
218 self.compile_function(func_idx, func_def)?;
219 func_idx += 1;
220 }
221
222 ItemKind::Impl(impl_block) => {
223 for method in &impl_block.methods {
224 self.compile_function(func_idx, method)?;
225 func_idx += 1;
226 }
227 }
228
229 _ => {}
230 }
231 }
232
233 {
234 let fake_func = crate::ast::FunctionDef {
235 name: "__script".to_string(),
236 type_params: vec![],
237 trait_bounds: vec![],
238 params: vec![],
239 return_type: None,
240 body: script_stmts,
241 is_method: false,
242 visibility: crate::ast::Visibility::Private,
243 };
244 let prev = self.current_module.clone();
245 if let Some(entry) = &self.entry_module {
246 self.current_module = Some(entry.clone());
247 }
248
249 self.compile_function(script_func_idx, &fake_func)?;
250 self.current_module = prev;
251 }
252
253 Ok(self.functions.clone())
254 }
255
256 pub(super) fn collect_top_level_locals(stmts: &[Stmt]) -> Vec<String> {
257 let mut names = Vec::new();
258 for stmt in stmts {
259 if let StmtKind::Local { bindings, .. } = &stmt.kind {
260 for binding in bindings {
261 names.push(binding.name.clone());
262 }
263 }
264 }
265
266 names
267 }
268
269 pub(super) fn is_module_init_function(
270 func_def: &crate::ast::FunctionDef,
271 module: &str,
272 ) -> bool {
273 func_def.name == format!("__init@{}", module)
274 }
275
276 pub(super) fn module_scope_name(&self) -> Option<&str> {
277 if let Some(module) = &self.current_module {
278 Some(module.as_str())
279 } else if let Some(entry) = &self.entry_module {
280 Some(entry.as_str())
281 } else {
282 Some("__root")
283 }
284 }
285
286 pub(super) fn looks_like_type_name(name: &str) -> bool {
287 name.chars()
288 .next()
289 .map(|c| c.is_ascii_uppercase())
290 .unwrap_or(false)
291 }
292
293 pub(super) fn is_initializer_context(&self) -> bool {
294 if let Some(name) = self.current_function_name.as_deref() {
295 name == "__script" || name.starts_with("__init@")
296 } else {
297 false
298 }
299 }
300
301 pub(super) fn current_scope_is_top_level(&self) -> bool {
302 matches!(self.scopes.last(), Some(scope) if scope.depth == 0)
303 }
304
305 pub(super) fn is_module_level_identifier(&self, name: &str) -> bool {
306 if let Some(module) = self.module_scope_name() {
307 if let Some(locals) = self.module_locals.get(module) {
308 return locals.contains(name);
309 }
310 }
311
312 false
313 }
314
315 pub(super) fn should_sync_module_local(&self, name: &str) -> bool {
316 self.is_initializer_context()
317 && self.current_scope_is_top_level()
318 && self.is_module_level_identifier(name)
319 }
320
321 pub(super) fn emit_store_module_global(&mut self, name: &str, src_reg: Register) {
322 if let Some(module) = self.module_scope_name() {
323 let key = format!("{}::{}", module, name);
324 let name_idx = self.add_string_constant(&key);
325 self.emit(Instruction::StoreGlobal(name_idx, src_reg), 0);
326 }
327 }
328
329 pub(super) fn emit_load_module_global(&mut self, name: &str, dest_reg: Register) -> Result<()> {
330 if let Some(module) = self.module_scope_name() {
331 let key = format!("{}::{}", module, name);
332 let const_idx = self.add_string_constant(&key);
333 self.emit(Instruction::LoadGlobal(dest_reg, const_idx), 0);
334 Ok(())
335 } else {
336 Err(LustError::CompileError(format!(
337 "Undefined variable: {}",
338 name
339 )))
340 }
341 }
342
343 pub(super) fn compile_function(
344 &mut self,
345 func_idx: usize,
346 func_def: &crate::ast::FunctionDef,
347 ) -> Result<()> {
348 self.current_function = func_idx;
349 self.next_register = 0;
350 self.max_register = 0;
351 self.scopes.clear();
352 let prev_function_name = self.current_function_name.clone();
353 self.current_function_name = Some(func_def.name.clone());
354 let mut scope = Scope {
355 locals: HashMap::new(),
356 depth: 0,
357 };
358 if func_def.is_method {
359 scope.locals.insert("self".to_string(), (0, false));
360 self.next_register = 1;
361 }
362
363 for param in &func_def.params {
364 let reg = self.allocate_register();
365 scope.locals.insert(param.name.clone(), (reg, false));
366 self.register_type(reg, param.ty.kind.clone());
367 }
368
369 self.scopes.push(scope);
370 for stmt in &func_def.body {
371 self.compile_stmt(stmt)?;
372 self.reset_temp_registers();
373 }
374
375 let last_instr = self.current_chunk().instructions.last();
376 if !matches!(last_instr, Some(Instruction::Return(_))) {
377 self.emit(Instruction::Return(255), 0);
378 }
379
380 self.functions[func_idx].set_register_count(self.max_register + 1);
381 self.scopes.pop();
382 self.current_function_name = prev_function_name;
383 Ok(())
384 }
385}