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