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