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