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