1use crate::{
2 ast::{FunctionDef, Item, ItemKind, UseTree, Visibility},
3 error::{LustError, Result},
4 lexer::Lexer,
5 parser::Parser,
6};
7use std::{
8 collections::{HashMap, HashSet},
9 fs,
10 path::{Path, PathBuf},
11};
12#[derive(Debug, Clone, Default)]
13pub struct ModuleImports {
14 pub function_aliases: HashMap<String, String>,
15 pub module_aliases: HashMap<String, String>,
16 pub type_aliases: HashMap<String, String>,
17}
18
19#[derive(Debug, Clone, Default)]
20pub struct ModuleExports {
21 pub functions: HashMap<String, String>,
22 pub types: HashMap<String, String>,
23}
24
25#[derive(Debug, Clone)]
26pub struct LoadedModule {
27 pub path: String,
28 pub items: Vec<Item>,
29 pub imports: ModuleImports,
30 pub exports: ModuleExports,
31 pub init_function: Option<String>,
32 pub source_path: PathBuf,
33}
34
35#[derive(Debug, Clone)]
36pub struct Program {
37 pub modules: Vec<LoadedModule>,
38 pub entry_module: String,
39}
40
41#[derive(Clone, Copy, Debug, Default)]
42struct ImportResolution {
43 import_value: bool,
44 import_type: bool,
45}
46
47impl ImportResolution {
48 fn both() -> Self {
49 Self {
50 import_value: true,
51 import_type: true,
52 }
53 }
54}
55
56pub struct ModuleLoader {
57 base_dir: PathBuf,
58 cache: HashMap<String, LoadedModule>,
59 visited: HashSet<String>,
60 source_overrides: HashMap<PathBuf, String>,
61}
62
63impl ModuleLoader {
64 pub fn new(base_dir: impl Into<PathBuf>) -> Self {
65 Self {
66 base_dir: base_dir.into(),
67 cache: HashMap::new(),
68 visited: HashSet::new(),
69 source_overrides: HashMap::new(),
70 }
71 }
72
73 pub fn set_source_overrides(&mut self, overrides: HashMap<PathBuf, String>) {
74 self.source_overrides = overrides;
75 }
76
77 pub fn set_source_override<P: Into<PathBuf>, S: Into<String>>(&mut self, path: P, source: S) {
78 self.source_overrides.insert(path.into(), source.into());
79 }
80
81 pub fn clear_source_overrides(&mut self) {
82 self.source_overrides.clear();
83 }
84
85 pub fn load_program_from_entry(&mut self, entry_file: &str) -> Result<Program> {
86 let entry_path = Path::new(entry_file);
87 let entry_dir = entry_path.parent().unwrap_or_else(|| Path::new("."));
88 self.base_dir = entry_dir.to_path_buf();
89 let entry_module = Self::module_path_for_file(entry_path);
90 let mut order: Vec<String> = Vec::new();
91 let mut stack: HashSet<String> = HashSet::new();
92 self.load_module_recursive(&entry_module, &mut order, &mut stack, true)?;
93 let modules = order
94 .into_iter()
95 .filter_map(|m| self.cache.get(&m).cloned())
96 .collect::<Vec<_>>();
97 Ok(Program {
98 modules,
99 entry_module,
100 })
101 }
102
103 fn load_module_recursive(
104 &mut self,
105 module_path: &str,
106 order: &mut Vec<String>,
107 stack: &mut HashSet<String>,
108 is_entry: bool,
109 ) -> Result<()> {
110 if self.visited.contains(module_path) {
111 return Ok(());
112 }
113
114 if !stack.insert(module_path.to_string()) {
115 return Ok(());
116 }
117
118 let mut loaded = self.load_single_module(module_path, is_entry)?;
119 self.cache.insert(module_path.to_string(), loaded.clone());
120 let deps = self.collect_dependencies(&loaded.items);
121 for dep in deps {
122 self.load_module_recursive(&dep, order, stack, false)?;
123 }
124
125 self.finalize_module(&mut loaded)?;
126 self.cache.insert(module_path.to_string(), loaded.clone());
127 self.visited.insert(module_path.to_string());
128 order.push(module_path.to_string());
129 stack.remove(module_path);
130 Ok(())
131 }
132
133 fn load_single_module(&self, module_path: &str, is_entry: bool) -> Result<LoadedModule> {
134 let file = self.file_for_module_path(module_path);
135 let source = if let Some(src) = self.source_overrides.get(&file) {
136 src.clone()
137 } else {
138 fs::read_to_string(&file).map_err(|e| {
139 LustError::Unknown(format!("Failed to read module '{}': {}", file.display(), e))
140 })?
141 };
142 let mut lexer = Lexer::new(&source);
143 let tokens = lexer
144 .tokenize()
145 .map_err(|err| Self::attach_module_to_error(err, module_path))?;
146 let mut parser = Parser::new(tokens);
147 let mut items = parser
148 .parse()
149 .map_err(|err| Self::attach_module_to_error(err, module_path))?;
150 let mut imports = ModuleImports::default();
151 let mut exports = ModuleExports::default();
152 let mut new_items: Vec<Item> = Vec::new();
153 let mut init_function: Option<String> = None;
154 for item in items.drain(..) {
155 match &item.kind {
156 ItemKind::Function(func) => {
157 let mut f = func.clone();
158 if !f.is_method && !f.name.contains(':') && !f.name.contains('.') {
159 let fq = format!("{}.{}", module_path, f.name);
160 imports.function_aliases.insert(f.name.clone(), fq.clone());
161 f.name = fq.clone();
162 if matches!(f.visibility, Visibility::Public) {
163 exports
164 .functions
165 .insert(self.simple_name(&f.name).to_string(), f.name.clone());
166 }
167 } else {
168 if matches!(f.visibility, Visibility::Public) {
169 exports
170 .functions
171 .insert(self.simple_name(&f.name).to_string(), f.name.clone());
172 }
173 }
174
175 new_items.push(Item::new(ItemKind::Function(f), item.span));
176 }
177
178 ItemKind::Struct(s) => {
179 if matches!(s.visibility, Visibility::Public) {
180 exports
181 .types
182 .insert(s.name.clone(), format!("{}.{}", module_path, s.name));
183 }
184
185 new_items.push(item);
186 }
187
188 ItemKind::Enum(e) => {
189 if matches!(e.visibility, Visibility::Public) {
190 exports
191 .types
192 .insert(e.name.clone(), format!("{}.{}", module_path, e.name));
193 }
194
195 new_items.push(item);
196 }
197
198 ItemKind::Trait(t) => {
199 if matches!(t.visibility, Visibility::Public) {
200 exports
201 .types
202 .insert(t.name.clone(), format!("{}.{}", module_path, t.name));
203 }
204
205 new_items.push(item);
206 }
207
208 ItemKind::TypeAlias { name, .. } => {
209 exports
210 .types
211 .insert(name.clone(), format!("{}.{}", module_path, name));
212 new_items.push(item);
213 }
214
215 ItemKind::Script(stmts) => {
216 if is_entry {
217 new_items.push(Item::new(ItemKind::Script(stmts.clone()), item.span));
218 } else {
219 let init_name = format!("__init@{}", module_path);
220 let func = FunctionDef {
221 name: init_name.clone(),
222 type_params: vec![],
223 trait_bounds: vec![],
224 params: vec![],
225 return_type: None,
226 body: stmts.clone(),
227 is_method: false,
228 visibility: Visibility::Private,
229 };
230 new_items.push(Item::new(ItemKind::Function(func), item.span));
231 init_function = Some(init_name);
232 }
233 }
234
235 ItemKind::Extern {
236 abi,
237 items: extern_items,
238 } => {
239 let mut rewritten = Vec::new();
240 for extern_item in extern_items {
241 match extern_item {
242 crate::ast::ExternItem::Function {
243 name,
244 params,
245 return_type,
246 } => {
247 let mut new_name = name.clone();
248 if !new_name.contains('.') && !new_name.contains(':') {
249 new_name = format!("{}.{}", module_path, new_name);
250 }
251
252 exports.functions.insert(
253 self.simple_name(&new_name).to_string(),
254 new_name.clone(),
255 );
256 imports.function_aliases.insert(
257 self.simple_name(&new_name).to_string(),
258 new_name.clone(),
259 );
260
261 rewritten.push(crate::ast::ExternItem::Function {
262 name: new_name,
263 params: params.clone(),
264 return_type: return_type.clone(),
265 });
266 }
267 }
268 }
269 new_items.push(Item::new(
270 ItemKind::Extern {
271 abi: abi.clone(),
272 items: rewritten,
273 },
274 item.span,
275 ));
276 }
277
278 _ => {
279 new_items.push(item);
280 }
281 }
282 }
283
284 Ok(LoadedModule {
285 path: module_path.to_string(),
286 items: new_items,
287 imports,
288 exports,
289 init_function,
290 source_path: file,
291 })
292 }
293
294 fn collect_dependencies(&self, items: &[Item]) -> Vec<String> {
295 let mut deps = HashSet::new();
296 for item in items {
297 if let ItemKind::Use { public: _, tree } = &item.kind {
298 self.collect_deps_from_use(tree, &mut deps);
299 }
300 }
301
302 deps.into_iter().collect()
303 }
304
305 fn finalize_module(&mut self, module: &mut LoadedModule) -> Result<()> {
306 for item in &module.items {
307 if let ItemKind::Use { tree, .. } = &item.kind {
308 self.process_use_tree(tree, &mut module.imports)?;
309 }
310 }
311
312 for item in &module.items {
313 if let ItemKind::Use { public: true, tree } = &item.kind {
314 self.apply_reexport(tree, &mut module.exports)?;
315 }
316 }
317
318 module
319 .imports
320 .module_aliases
321 .entry(self.simple_tail(&module.path).to_string())
322 .or_insert_with(|| module.path.clone());
323 Ok(())
324 }
325
326 fn collect_deps_from_use(&self, tree: &UseTree, deps: &mut HashSet<String>) {
327 match tree {
328 UseTree::Path {
329 path,
330 alias: _,
331 import_module: _,
332 } => {
333 let full = path.join(".");
334 let full_file = self.file_for_module_path(&full);
335 if self.module_source_known(&full, &full_file) {
336 deps.insert(full);
337 } else if path.len() > 1 {
338 deps.insert(path[..path.len() - 1].join("."));
339 }
340 }
341
342 UseTree::Group { prefix, items } => {
343 let module = prefix.join(".");
344 if !module.is_empty() {
345 deps.insert(module);
346 }
347
348 for item in items {
349 if item.path.len() > 1 {
350 let mut combined: Vec<String> = prefix.clone();
351 combined.extend(item.path[..item.path.len() - 1].iter().cloned());
352 let module_path = combined.join(".");
353 if !module_path.is_empty() {
354 deps.insert(module_path);
355 }
356 }
357 }
358 }
359
360 UseTree::Glob { prefix } => {
361 deps.insert(prefix.join("."));
362 }
363 }
364 }
365
366 fn process_use_tree(&self, tree: &UseTree, imports: &mut ModuleImports) -> Result<()> {
367 match tree {
368 UseTree::Path { path, alias, .. } => {
369 let full = path.join(".");
370 let full_file = self.file_for_module_path(&full);
371 if self.module_source_known(&full, &full_file) {
372 let alias_name = alias
373 .clone()
374 .unwrap_or_else(|| path.last().unwrap().clone());
375 imports.module_aliases.insert(alias_name, full);
376 } else if path.len() > 1 {
377 let module = path[..path.len() - 1].join(".");
378 let item = path.last().unwrap().clone();
379 let alias_name = alias.clone().unwrap_or_else(|| item.clone());
380 let classification = self.classify_import_target(&module, &item);
381 let fq = format!("{}.{}", module, item);
382 if classification.import_value {
383 imports
384 .function_aliases
385 .insert(alias_name.clone(), fq.clone());
386 }
387
388 if classification.import_type {
389 imports.type_aliases.insert(alias_name, fq);
390 }
391 }
392 }
393
394 UseTree::Group { prefix, items } => {
395 for item in items {
396 if item.path.is_empty() {
397 continue;
398 }
399
400 let alias_name = item
401 .alias
402 .clone()
403 .unwrap_or_else(|| item.path.last().unwrap().clone());
404 let mut full_segments = prefix.clone();
405 full_segments.extend(item.path.clone());
406 let full = full_segments.join(".");
407 let full_file = self.file_for_module_path(&full);
408 if self.module_source_known(&full, &full_file) {
409 imports.module_aliases.insert(alias_name, full);
410 continue;
411 }
412
413 let mut module_segments = full_segments.clone();
414 let item_name = module_segments.pop().unwrap();
415 let module_path = module_segments.join(".");
416 let fq_name = if module_path.is_empty() {
417 item_name.clone()
418 } else {
419 format!("{}.{}", module_path, item_name)
420 };
421 let classification = self.classify_import_target(&module_path, &item_name);
422 if classification.import_value {
423 imports
424 .function_aliases
425 .insert(alias_name.clone(), fq_name.clone());
426 }
427
428 if classification.import_type {
429 imports.type_aliases.insert(alias_name.clone(), fq_name);
430 }
431 }
432 }
433
434 UseTree::Glob { prefix } => {
435 let module = prefix.join(".");
436 if let Some(loaded) = self.cache.get(&module) {
437 for (name, fq) in &loaded.exports.functions {
438 imports.function_aliases.insert(name.clone(), fq.clone());
439 }
440
441 for (name, fq) in &loaded.exports.types {
442 imports.type_aliases.insert(name.clone(), fq.clone());
443 }
444 }
445
446 let alias_name = prefix.last().cloned().unwrap_or_else(|| module.clone());
447 if !module.is_empty() {
448 imports.module_aliases.insert(alias_name, module);
449 }
450 }
451 }
452
453 Ok(())
454 }
455
456 fn attach_module_to_error(error: LustError, module_path: &str) -> LustError {
457 match error {
458 LustError::LexerError {
459 line,
460 column,
461 message,
462 module,
463 } => LustError::LexerError {
464 line,
465 column,
466 message,
467 module: module.or_else(|| Some(module_path.to_string())),
468 },
469 LustError::ParserError {
470 line,
471 column,
472 message,
473 module,
474 } => LustError::ParserError {
475 line,
476 column,
477 message,
478 module: module.or_else(|| Some(module_path.to_string())),
479 },
480 LustError::CompileErrorWithSpan {
481 message,
482 line,
483 column,
484 module,
485 } => LustError::CompileErrorWithSpan {
486 message,
487 line,
488 column,
489 module: module.or_else(|| Some(module_path.to_string())),
490 },
491 other => other,
492 }
493 }
494
495 fn apply_reexport(&self, tree: &UseTree, exports: &mut ModuleExports) -> Result<()> {
496 match tree {
497 UseTree::Path { path, alias, .. } => {
498 if path.len() == 1 {
499 return Ok(());
500 }
501
502 let module = path[..path.len() - 1].join(".");
503 let item = path.last().unwrap().clone();
504 let alias_name = alias.clone().unwrap_or_else(|| item.clone());
505 let fq = format!("{}.{}", module, item);
506 let classification = self.classify_import_target(&module, &item);
507 if classification.import_type {
508 exports.types.insert(alias_name.clone(), fq.clone());
509 }
510
511 if classification.import_value {
512 exports.functions.insert(alias_name, fq);
513 }
514
515 Ok(())
516 }
517
518 UseTree::Group { prefix, items } => {
519 for item in items {
520 if item.path.is_empty() {
521 continue;
522 }
523
524 let alias_name = item
525 .alias
526 .clone()
527 .unwrap_or_else(|| item.path.last().unwrap().clone());
528 let mut full_segments = prefix.clone();
529 full_segments.extend(item.path.clone());
530 let full = full_segments.join(".");
531 let full_file = self.file_for_module_path(&full);
532 if self.module_source_known(&full, &full_file) {
533 continue;
534 }
535
536 let mut module_segments = full_segments.clone();
537 let item_name = module_segments.pop().unwrap();
538 let module_path = module_segments.join(".");
539 let fq_name = if module_path.is_empty() {
540 item_name.clone()
541 } else {
542 format!("{}.{}", module_path, item_name)
543 };
544 let classification = self.classify_import_target(&module_path, &item_name);
545 if classification.import_type {
546 exports.types.insert(alias_name.clone(), fq_name.clone());
547 }
548
549 if classification.import_value {
550 exports.functions.insert(alias_name.clone(), fq_name);
551 }
552 }
553
554 Ok(())
555 }
556
557 UseTree::Glob { prefix } => {
558 let module = prefix.join(".");
559 if let Some(loaded) = self.cache.get(&module) {
560 for (n, fq) in &loaded.exports.types {
561 exports.types.insert(n.clone(), fq.clone());
562 }
563
564 for (n, fq) in &loaded.exports.functions {
565 exports.functions.insert(n.clone(), fq.clone());
566 }
567 }
568
569 Ok(())
570 }
571 }
572 }
573
574 fn simple_name<'a>(&self, qualified: &'a str) -> &'a str {
575 qualified
576 .rsplit_once('.')
577 .map(|(_, n)| n)
578 .unwrap_or(qualified)
579 }
580
581 fn simple_tail<'a>(&self, module_path: &'a str) -> &'a str {
582 module_path
583 .rsplit_once('.')
584 .map(|(_, n)| n)
585 .unwrap_or(module_path)
586 }
587
588 fn module_source_known(&self, module_path: &str, file: &Path) -> bool {
589 file.exists()
590 || self.source_overrides.contains_key(file)
591 || self.cache.contains_key(module_path)
592 }
593
594 fn classify_import_target(&self, module_path: &str, item_name: &str) -> ImportResolution {
595 if module_path.is_empty() {
596 return ImportResolution::both();
597 }
598
599 if let Some(module) = self.cache.get(module_path) {
600 let has_value = module.exports.functions.contains_key(item_name);
601 let has_type = module.exports.types.contains_key(item_name);
602 if has_value || has_type {
603 return ImportResolution {
604 import_value: has_value,
605 import_type: has_type,
606 };
607 }
608 }
609
610 ImportResolution::both()
611 }
612
613 fn file_for_module_path(&self, module_path: &str) -> PathBuf {
614 let mut p = self.base_dir.clone();
615 for seg in module_path.split('.') {
616 p.push(seg);
617 }
618
619 p.set_extension("lust");
620 p
621 }
622
623 fn module_path_for_file(path: &Path) -> String {
624 let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("");
625 stem.to_string()
626 }
627}