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
257 rewritten.push(crate::ast::ExternItem::Function {
258 name: new_name,
259 params: params.clone(),
260 return_type: return_type.clone(),
261 });
262 }
263 }
264 }
265 new_items.push(Item::new(
266 ItemKind::Extern {
267 abi: abi.clone(),
268 items: rewritten,
269 },
270 item.span,
271 ));
272 }
273
274 _ => {
275 new_items.push(item);
276 }
277 }
278 }
279
280 Ok(LoadedModule {
281 path: module_path.to_string(),
282 items: new_items,
283 imports,
284 exports,
285 init_function,
286 source_path: file,
287 })
288 }
289
290 fn collect_dependencies(&self, items: &[Item]) -> Vec<String> {
291 let mut deps = HashSet::new();
292 for item in items {
293 if let ItemKind::Use { public: _, tree } = &item.kind {
294 self.collect_deps_from_use(tree, &mut deps);
295 }
296 }
297
298 deps.into_iter().collect()
299 }
300
301 fn finalize_module(&mut self, module: &mut LoadedModule) -> Result<()> {
302 for item in &module.items {
303 if let ItemKind::Use { tree, .. } = &item.kind {
304 self.process_use_tree(tree, &mut module.imports)?;
305 }
306 }
307
308 for item in &module.items {
309 if let ItemKind::Use { public: true, tree } = &item.kind {
310 self.apply_reexport(tree, &mut module.exports)?;
311 }
312 }
313
314 module
315 .imports
316 .module_aliases
317 .entry(self.simple_tail(&module.path).to_string())
318 .or_insert_with(|| module.path.clone());
319 Ok(())
320 }
321
322 fn collect_deps_from_use(&self, tree: &UseTree, deps: &mut HashSet<String>) {
323 match tree {
324 UseTree::Path {
325 path,
326 alias: _,
327 import_module: _,
328 } => {
329 let full = path.join(".");
330 let full_file = self.file_for_module_path(&full);
331 if self.module_source_known(&full, &full_file) {
332 deps.insert(full);
333 } else if path.len() > 1 {
334 deps.insert(path[..path.len() - 1].join("."));
335 }
336 }
337
338 UseTree::Group { prefix, items } => {
339 let module = prefix.join(".");
340 if !module.is_empty() {
341 deps.insert(module);
342 }
343
344 for item in items {
345 if item.path.len() > 1 {
346 let mut combined: Vec<String> = prefix.clone();
347 combined.extend(item.path[..item.path.len() - 1].iter().cloned());
348 let module_path = combined.join(".");
349 if !module_path.is_empty() {
350 deps.insert(module_path);
351 }
352 }
353 }
354 }
355
356 UseTree::Glob { prefix } => {
357 deps.insert(prefix.join("."));
358 }
359 }
360 }
361
362 fn process_use_tree(&self, tree: &UseTree, imports: &mut ModuleImports) -> Result<()> {
363 match tree {
364 UseTree::Path { path, alias, .. } => {
365 let full = path.join(".");
366 let full_file = self.file_for_module_path(&full);
367 if self.module_source_known(&full, &full_file) {
368 let alias_name = alias
369 .clone()
370 .unwrap_or_else(|| path.last().unwrap().clone());
371 imports.module_aliases.insert(alias_name, full);
372 } else if path.len() > 1 {
373 let module = path[..path.len() - 1].join(".");
374 let item = path.last().unwrap().clone();
375 let alias_name = alias.clone().unwrap_or_else(|| item.clone());
376 let classification = self.classify_import_target(&module, &item);
377 let fq = format!("{}.{}", module, item);
378 if classification.import_value {
379 imports
380 .function_aliases
381 .insert(alias_name.clone(), fq.clone());
382 }
383
384 if classification.import_type {
385 imports.type_aliases.insert(alias_name, fq);
386 }
387 }
388 }
389
390 UseTree::Group { prefix, items } => {
391 for item in items {
392 if item.path.is_empty() {
393 continue;
394 }
395
396 let alias_name = item
397 .alias
398 .clone()
399 .unwrap_or_else(|| item.path.last().unwrap().clone());
400 let mut full_segments = prefix.clone();
401 full_segments.extend(item.path.clone());
402 let full = full_segments.join(".");
403 let full_file = self.file_for_module_path(&full);
404 if self.module_source_known(&full, &full_file) {
405 imports.module_aliases.insert(alias_name, full);
406 continue;
407 }
408
409 let mut module_segments = full_segments.clone();
410 let item_name = module_segments.pop().unwrap();
411 let module_path = module_segments.join(".");
412 let fq_name = if module_path.is_empty() {
413 item_name.clone()
414 } else {
415 format!("{}.{}", module_path, item_name)
416 };
417 let classification = self.classify_import_target(&module_path, &item_name);
418 if classification.import_value {
419 imports
420 .function_aliases
421 .insert(alias_name.clone(), fq_name.clone());
422 }
423
424 if classification.import_type {
425 imports.type_aliases.insert(alias_name.clone(), fq_name);
426 }
427 }
428 }
429
430 UseTree::Glob { prefix } => {
431 let module = prefix.join(".");
432 if let Some(loaded) = self.cache.get(&module) {
433 for (name, fq) in &loaded.exports.functions {
434 imports.function_aliases.insert(name.clone(), fq.clone());
435 }
436
437 for (name, fq) in &loaded.exports.types {
438 imports.type_aliases.insert(name.clone(), fq.clone());
439 }
440 }
441
442 let alias_name = prefix.last().cloned().unwrap_or_else(|| module.clone());
443 if !module.is_empty() {
444 imports.module_aliases.insert(alias_name, module);
445 }
446 }
447 }
448
449 Ok(())
450 }
451
452 fn attach_module_to_error(error: LustError, module_path: &str) -> LustError {
453 match error {
454 LustError::LexerError {
455 line,
456 column,
457 message,
458 module,
459 } => LustError::LexerError {
460 line,
461 column,
462 message,
463 module: module.or_else(|| Some(module_path.to_string())),
464 },
465 LustError::ParserError {
466 line,
467 column,
468 message,
469 module,
470 } => LustError::ParserError {
471 line,
472 column,
473 message,
474 module: module.or_else(|| Some(module_path.to_string())),
475 },
476 LustError::CompileErrorWithSpan {
477 message,
478 line,
479 column,
480 module,
481 } => LustError::CompileErrorWithSpan {
482 message,
483 line,
484 column,
485 module: module.or_else(|| Some(module_path.to_string())),
486 },
487 other => other,
488 }
489 }
490
491 fn apply_reexport(&self, tree: &UseTree, exports: &mut ModuleExports) -> Result<()> {
492 match tree {
493 UseTree::Path { path, alias, .. } => {
494 if path.len() == 1 {
495 return Ok(());
496 }
497
498 let module = path[..path.len() - 1].join(".");
499 let item = path.last().unwrap().clone();
500 let alias_name = alias.clone().unwrap_or_else(|| item.clone());
501 let fq = format!("{}.{}", module, item);
502 let classification = self.classify_import_target(&module, &item);
503 if classification.import_type {
504 exports.types.insert(alias_name.clone(), fq.clone());
505 }
506
507 if classification.import_value {
508 exports.functions.insert(alias_name, fq);
509 }
510
511 Ok(())
512 }
513
514 UseTree::Group { prefix, items } => {
515 for item in items {
516 if item.path.is_empty() {
517 continue;
518 }
519
520 let alias_name = item
521 .alias
522 .clone()
523 .unwrap_or_else(|| item.path.last().unwrap().clone());
524 let mut full_segments = prefix.clone();
525 full_segments.extend(item.path.clone());
526 let full = full_segments.join(".");
527 let full_file = self.file_for_module_path(&full);
528 if self.module_source_known(&full, &full_file) {
529 continue;
530 }
531
532 let mut module_segments = full_segments.clone();
533 let item_name = module_segments.pop().unwrap();
534 let module_path = module_segments.join(".");
535 let fq_name = if module_path.is_empty() {
536 item_name.clone()
537 } else {
538 format!("{}.{}", module_path, item_name)
539 };
540 let classification = self.classify_import_target(&module_path, &item_name);
541 if classification.import_type {
542 exports.types.insert(alias_name.clone(), fq_name.clone());
543 }
544
545 if classification.import_value {
546 exports.functions.insert(alias_name.clone(), fq_name);
547 }
548 }
549
550 Ok(())
551 }
552
553 UseTree::Glob { prefix } => {
554 let module = prefix.join(".");
555 if let Some(loaded) = self.cache.get(&module) {
556 for (n, fq) in &loaded.exports.types {
557 exports.types.insert(n.clone(), fq.clone());
558 }
559
560 for (n, fq) in &loaded.exports.functions {
561 exports.functions.insert(n.clone(), fq.clone());
562 }
563 }
564
565 Ok(())
566 }
567 }
568 }
569
570 fn simple_name<'a>(&self, qualified: &'a str) -> &'a str {
571 qualified
572 .rsplit_once('.')
573 .map(|(_, n)| n)
574 .unwrap_or(qualified)
575 }
576
577 fn simple_tail<'a>(&self, module_path: &'a str) -> &'a str {
578 module_path
579 .rsplit_once('.')
580 .map(|(_, n)| n)
581 .unwrap_or(module_path)
582 }
583
584 fn module_source_known(&self, module_path: &str, file: &Path) -> bool {
585 file.exists()
586 || self.source_overrides.contains_key(file)
587 || self.cache.contains_key(module_path)
588 }
589
590 fn classify_import_target(&self, module_path: &str, item_name: &str) -> ImportResolution {
591 if module_path.is_empty() {
592 return ImportResolution::both();
593 }
594
595 if let Some(module) = self.cache.get(module_path) {
596 let has_value = module.exports.functions.contains_key(item_name);
597 let has_type = module.exports.types.contains_key(item_name);
598 if has_value || has_type {
599 return ImportResolution {
600 import_value: has_value,
601 import_type: has_type,
602 };
603 }
604 }
605
606 ImportResolution::both()
607 }
608
609 fn file_for_module_path(&self, module_path: &str) -> PathBuf {
610 let mut p = self.base_dir.clone();
611 for seg in module_path.split('.') {
612 p.push(seg);
613 }
614
615 p.set_extension("lust");
616 p
617 }
618
619 fn module_path_for_file(path: &Path) -> String {
620 let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("");
621 stem.to_string()
622 }
623}