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, Vec<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
109 .entry(prefix.into())
110 .or_default()
111 .push(ModuleRoot {
112 base: root.into(),
113 root_module,
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 let Some((head, tail)) = new_name.split_once(':') {
281 let qualified_head =
282 if head.contains('.') || head.contains("::") {
283 head.to_string()
284 } else {
285 format!("{}.{}", module_path, head)
286 };
287 new_name = format!("{}:{}", qualified_head, tail);
288 } else if !new_name.contains('.') {
289 new_name = format!("{}.{}", module_path, new_name);
290 }
291
292 exports.functions.insert(
293 self.simple_name(&new_name).to_string(),
294 new_name.clone(),
295 );
296 imports.function_aliases.insert(
297 self.simple_name(&new_name).to_string(),
298 new_name.clone(),
299 );
300
301 rewritten.push(crate::ast::ExternItem::Function {
302 name: new_name,
303 params: params.clone(),
304 return_type: return_type.clone(),
305 });
306 }
307 }
308 }
309 new_items.push(Item::new(
310 ItemKind::Extern {
311 abi: abi.clone(),
312 items: rewritten,
313 },
314 item.span,
315 ));
316 }
317
318 _ => {
319 new_items.push(item);
320 }
321 }
322 }
323
324 Ok(LoadedModule {
325 path: module_path.to_string(),
326 items: new_items,
327 imports,
328 exports,
329 init_function,
330 source_path: file,
331 })
332 }
333
334 fn collect_dependencies(&self, items: &[Item]) -> Vec<String> {
335 let mut deps = HashSet::new();
336 for item in items {
337 if let ItemKind::Use { public: _, tree } = &item.kind {
338 self.collect_deps_from_use(tree, &mut deps);
339 }
340 }
341
342 deps.into_iter().collect()
343 }
344
345 fn finalize_module(&mut self, module: &mut LoadedModule) -> Result<()> {
346 for item in &module.items {
347 if let ItemKind::Use { tree, .. } = &item.kind {
348 self.process_use_tree(tree, &mut module.imports)?;
349 }
350 }
351
352 for item in &module.items {
353 if let ItemKind::Use { public: true, tree } = &item.kind {
354 self.apply_reexport(tree, &mut module.exports)?;
355 }
356 }
357
358 module
359 .imports
360 .module_aliases
361 .entry(self.simple_tail(&module.path).to_string())
362 .or_insert_with(|| module.path.clone());
363 Ok(())
364 }
365
366 fn collect_deps_from_use(&self, tree: &UseTree, deps: &mut HashSet<String>) {
367 match tree {
368 UseTree::Path {
369 path,
370 alias: _,
371 import_module: _,
372 } => {
373 let full = path.join(".");
374 let full_file = self.file_for_module_path(&full);
375 if self.module_source_known(&full, &full_file) {
376 deps.insert(full);
377 } else if path.len() > 1 {
378 deps.insert(path[..path.len() - 1].join("."));
379 }
380 }
381
382 UseTree::Group { prefix, items } => {
383 let module = prefix.join(".");
384 if !module.is_empty() {
385 deps.insert(module);
386 }
387
388 for item in items {
389 if item.path.len() > 1 {
390 let mut combined: Vec<String> = prefix.clone();
391 combined.extend(item.path[..item.path.len() - 1].iter().cloned());
392 let module_path = combined.join(".");
393 if !module_path.is_empty() {
394 deps.insert(module_path);
395 }
396 }
397 }
398 }
399
400 UseTree::Glob { prefix } => {
401 deps.insert(prefix.join("."));
402 }
403 }
404 }
405
406 fn process_use_tree(&self, tree: &UseTree, imports: &mut ModuleImports) -> Result<()> {
407 match tree {
408 UseTree::Path { path, alias, .. } => {
409 let full = path.join(".");
410 let full_file = self.file_for_module_path(&full);
411 if self.module_source_known(&full, &full_file) {
412 let alias_name = alias
413 .clone()
414 .unwrap_or_else(|| path.last().unwrap().clone());
415 imports.module_aliases.insert(alias_name, full);
416 } else if path.len() > 1 {
417 let module = path[..path.len() - 1].join(".");
418 let item = path.last().unwrap().clone();
419 let alias_name = alias.clone().unwrap_or_else(|| item.clone());
420 let classification = self.classify_import_target(&module, &item);
421 let fq = format!("{}.{}", module, item);
422 if classification.import_value {
423 imports
424 .function_aliases
425 .insert(alias_name.clone(), fq.clone());
426 }
427
428 if classification.import_type {
429 imports.type_aliases.insert(alias_name, fq);
430 }
431 }
432 }
433
434 UseTree::Group { prefix, items } => {
435 for item in items {
436 if item.path.is_empty() {
437 continue;
438 }
439
440 let alias_name = item
441 .alias
442 .clone()
443 .unwrap_or_else(|| item.path.last().unwrap().clone());
444 let mut full_segments = prefix.clone();
445 full_segments.extend(item.path.clone());
446 let full = full_segments.join(".");
447 let full_file = self.file_for_module_path(&full);
448 if self.module_source_known(&full, &full_file) {
449 imports.module_aliases.insert(alias_name, full);
450 continue;
451 }
452
453 let mut module_segments = full_segments.clone();
454 let item_name = module_segments.pop().unwrap();
455 let module_path = module_segments.join(".");
456 let fq_name = if module_path.is_empty() {
457 item_name.clone()
458 } else {
459 format!("{}.{}", module_path, item_name)
460 };
461 let classification = self.classify_import_target(&module_path, &item_name);
462 if classification.import_value {
463 imports
464 .function_aliases
465 .insert(alias_name.clone(), fq_name.clone());
466 }
467
468 if classification.import_type {
469 imports.type_aliases.insert(alias_name.clone(), fq_name);
470 }
471 }
472 }
473
474 UseTree::Glob { prefix } => {
475 let module = prefix.join(".");
476 if let Some(loaded) = self.cache.get(&module) {
477 for (name, fq) in &loaded.exports.functions {
478 imports.function_aliases.insert(name.clone(), fq.clone());
479 }
480
481 for (name, fq) in &loaded.exports.types {
482 imports.type_aliases.insert(name.clone(), fq.clone());
483 }
484 }
485
486 let alias_name = prefix.last().cloned().unwrap_or_else(|| module.clone());
487 if !module.is_empty() {
488 imports.module_aliases.insert(alias_name, module);
489 }
490 }
491 }
492
493 Ok(())
494 }
495
496 fn attach_module_to_error(error: LustError, module_path: &str) -> LustError {
497 match error {
498 LustError::LexerError {
499 line,
500 column,
501 message,
502 module,
503 } => LustError::LexerError {
504 line,
505 column,
506 message,
507 module: module.or_else(|| Some(module_path.to_string())),
508 },
509 LustError::ParserError {
510 line,
511 column,
512 message,
513 module,
514 } => LustError::ParserError {
515 line,
516 column,
517 message,
518 module: module.or_else(|| Some(module_path.to_string())),
519 },
520 LustError::CompileErrorWithSpan {
521 message,
522 line,
523 column,
524 module,
525 } => LustError::CompileErrorWithSpan {
526 message,
527 line,
528 column,
529 module: module.or_else(|| Some(module_path.to_string())),
530 },
531 other => other,
532 }
533 }
534
535 fn apply_reexport(&self, tree: &UseTree, exports: &mut ModuleExports) -> Result<()> {
536 match tree {
537 UseTree::Path { path, alias, .. } => {
538 if path.len() == 1 {
539 return Ok(());
540 }
541
542 let module = path[..path.len() - 1].join(".");
543 let item = path.last().unwrap().clone();
544 let alias_name = alias.clone().unwrap_or_else(|| item.clone());
545 let fq = format!("{}.{}", module, item);
546 let classification = self.classify_import_target(&module, &item);
547 if classification.import_type {
548 exports.types.insert(alias_name.clone(), fq.clone());
549 }
550
551 if classification.import_value {
552 exports.functions.insert(alias_name, fq);
553 }
554
555 Ok(())
556 }
557
558 UseTree::Group { prefix, items } => {
559 for item in items {
560 if item.path.is_empty() {
561 continue;
562 }
563
564 let alias_name = item
565 .alias
566 .clone()
567 .unwrap_or_else(|| item.path.last().unwrap().clone());
568 let mut full_segments = prefix.clone();
569 full_segments.extend(item.path.clone());
570 let full = full_segments.join(".");
571 let full_file = self.file_for_module_path(&full);
572 if self.module_source_known(&full, &full_file) {
573 continue;
574 }
575
576 let mut module_segments = full_segments.clone();
577 let item_name = module_segments.pop().unwrap();
578 let module_path = module_segments.join(".");
579 let fq_name = if module_path.is_empty() {
580 item_name.clone()
581 } else {
582 format!("{}.{}", module_path, item_name)
583 };
584 let classification = self.classify_import_target(&module_path, &item_name);
585 if classification.import_type {
586 exports.types.insert(alias_name.clone(), fq_name.clone());
587 }
588
589 if classification.import_value {
590 exports.functions.insert(alias_name.clone(), fq_name);
591 }
592 }
593
594 Ok(())
595 }
596
597 UseTree::Glob { prefix } => {
598 let module = prefix.join(".");
599 if let Some(loaded) = self.cache.get(&module) {
600 for (n, fq) in &loaded.exports.types {
601 exports.types.insert(n.clone(), fq.clone());
602 }
603
604 for (n, fq) in &loaded.exports.functions {
605 exports.functions.insert(n.clone(), fq.clone());
606 }
607 }
608
609 Ok(())
610 }
611 }
612 }
613
614 fn simple_name<'a>(&self, qualified: &'a str) -> &'a str {
615 qualified
616 .rsplit_once('.')
617 .map(|(_, n)| n)
618 .unwrap_or(qualified)
619 }
620
621 fn simple_tail<'a>(&self, module_path: &'a str) -> &'a str {
622 module_path
623 .rsplit_once('.')
624 .map(|(_, n)| n)
625 .unwrap_or(module_path)
626 }
627
628 fn module_source_known(&self, module_path: &str, file: &Path) -> bool {
629 file.exists()
630 || self.source_overrides.contains_key(file)
631 || self.cache.contains_key(module_path)
632 }
633
634 fn classify_import_target(&self, module_path: &str, item_name: &str) -> ImportResolution {
635 if module_path.is_empty() {
636 return ImportResolution::both();
637 }
638
639 if let Some(module) = self.cache.get(module_path) {
640 let has_value = module.exports.functions.contains_key(item_name);
641 let has_type = module.exports.types.contains_key(item_name);
642 if has_value || has_type {
643 return ImportResolution {
644 import_value: has_value,
645 import_type: has_type,
646 };
647 }
648 }
649
650 ImportResolution::both()
651 }
652
653 fn file_for_module_path(&self, module_path: &str) -> PathBuf {
654 let segments: Vec<&str> = module_path.split('.').collect();
655 let candidates = self.resolve_dependency_roots(&segments);
656 if !candidates.is_empty() {
657 let mut fallback: Option<PathBuf> = None;
658 for (root, consumed) in candidates.iter().rev() {
659 let candidate = Self::path_from_root(root, &segments, *consumed);
660 if candidate.exists() {
661 return candidate;
662 }
663 if fallback.is_none() {
664 fallback = Some(candidate);
665 }
666 }
667 if let Some(path) = fallback {
668 return path;
669 }
670 }
671
672 let mut fallback = self.base_dir.clone();
673 for seg in &segments {
674 fallback.push(seg);
675 }
676 fallback.set_extension("lust");
677 fallback
678 }
679
680 fn resolve_dependency_roots(&self, segments: &[&str]) -> Vec<(&ModuleRoot, usize)> {
681 let mut matched: Vec<(&ModuleRoot, usize)> = Vec::new();
682 let mut prefix_segments: Vec<&str> = Vec::new();
683 for (index, segment) in segments.iter().enumerate() {
684 prefix_segments.push(*segment);
685 let key = prefix_segments.join(".");
686 if let Some(roots) = self.module_roots.get(&key) {
687 for root in roots {
688 matched.push((root, index + 1));
689 }
690 }
691 }
692 matched
693 }
694
695 fn path_from_root(root: &ModuleRoot, segments: &[&str], consumed: usize) -> PathBuf {
696 let mut path = root.base.clone();
697 if consumed == segments.len() {
698 if let Some(relative) = &root.root_module {
699 path.push(relative);
700 } else if let Some(last) = segments.last() {
701 path.push(format!("{last}.lust"));
702 }
703 return path;
704 }
705 for seg in &segments[consumed..segments.len() - 1] {
706 path.push(seg);
707 }
708 if let Some(last) = segments.last() {
709 path.push(format!("{last}.lust"));
710 }
711 path
712 }
713
714 fn module_path_for_file(path: &Path) -> String {
715 let stem = path.file_stem().and_then(|s| s.to_str()).unwrap_or("");
716 stem.to_string()
717 }
718}
719
720#[cfg(not(feature = "std"))]
721pub struct ModuleLoader {
722 cache: HashMap<String, LoadedModule>,
723 visited: HashSet<String>,
724}
725
726#[cfg(not(feature = "std"))]
727impl ModuleLoader {
728 pub fn new() -> Self {
729 Self {
730 cache: HashMap::new(),
731 visited: HashSet::new(),
732 }
733 }
734
735 pub fn clear_cache(&mut self) {
736 self.cache.clear();
737 self.visited.clear();
738 }
739
740 pub fn load_program_from_modules(
741 &mut self,
742 modules: Vec<LoadedModule>,
743 entry_module: String,
744 ) -> Program {
745 Program {
746 modules,
747 entry_module,
748 }
749 }
750}