use crate::ast::AstNode;
use crate::diagnostic::{ColorConfig, DiagnosticEmitter, Files, StandardEmitter, ToDiagnostic};
use crate::embedded_std::embedded_std_sources;
use crate::lexer::Lexer;
use crate::parser::Parser;
use crate::source::Source;
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
pub struct ModuleResolver {
base_path: PathBuf,
compiled_modules: HashMap<String, Vec<AstNode>>,
import_stack: Vec<String>,
being_imported: HashSet<String>,
canonical_cache: HashMap<PathBuf, String>, embedded_sources: &'static HashMap<String, String>,
}
impl ModuleResolver {
pub fn new(base_path: PathBuf) -> Self {
Self {
base_path,
compiled_modules: HashMap::new(),
import_stack: Vec::new(),
being_imported: HashSet::new(),
canonical_cache: HashMap::new(),
embedded_sources: embedded_std_sources(),
}
}
fn resolve_embedded_key<'a>(&'a self, module_path: &'a str) -> Option<&'a str> {
if self.embedded_sources.contains_key(module_path) {
return Some(module_path);
}
module_path
.strip_prefix("std.")
.filter(|short| self.embedded_sources.contains_key(*short))
}
fn normalize_module_key<'a>(&'a self, module_path: &'a str) -> &'a str {
self.resolve_embedded_key(module_path)
.unwrap_or(module_path)
}
pub fn has_embedded_module(&self, module_path: &str) -> bool {
self.resolve_embedded_key(module_path).is_some()
}
fn resolve_embedded_module_if_any(
&mut self,
module_path: &str,
files: &mut Files,
) -> Result<Option<Vec<AstNode>>, String> {
if let Some(embedded_key) = self.resolve_embedded_key(module_path) {
let embedded_key = embedded_key.to_string();
let cache_key = self.normalize_module_key(module_path).to_string();
if let Some(nodes) = self.compiled_modules.get(&cache_key) {
return Ok(Some(nodes.clone()));
}
if self.being_imported.contains(&cache_key) {
return Err(format!(
"Circular import detected: {} -> {}",
self.import_stack.join(" -> "),
cache_key
));
}
self.being_imported.insert(cache_key.clone());
self.import_stack.push(cache_key.clone());
let source = self
.embedded_sources
.get(&embedded_key)
.ok_or_else(|| format!("Embedded module not found: {}", module_path))?;
let virtual_path =
PathBuf::from(format!("<embedded>/{}.mux", embedded_key.replace('.', "/")));
let nodes = self.parse_module(&virtual_path, files, Some(source.as_str()))?;
return Ok(Some(nodes));
}
Ok(None)
}
fn determine_file_path(
&self,
module_path: &str,
current_file: Option<&Path>,
) -> Result<PathBuf, String> {
if module_path.starts_with("./") || module_path.starts_with("../") {
let current_dir = current_file
.and_then(|p| p.parent())
.ok_or("Cannot resolve relative import: no current file")?;
let relative_path = module_path.trim_start_matches("./");
let mut path = current_dir.to_path_buf();
for part in relative_path.split('/') {
if part == ".." {
path.pop();
} else if !part.is_empty() {
path.push(part);
}
}
path.set_extension("mux");
Ok(path)
} else if module_path.starts_with('/') {
let mut path = PathBuf::from(module_path);
path.set_extension("mux");
Ok(path)
} else {
self.module_path_to_file(module_path)
}
}
pub fn resolve_import_path(
&mut self,
module_path: &str,
current_file: Option<&Path>,
files: &mut Files,
) -> Result<Vec<AstNode>, String> {
if let Some(nodes) = self.resolve_embedded_module_if_any(module_path, files)? {
return Ok(nodes);
}
let file_path = self.determine_file_path(module_path, current_file)?;
let canonical_path = file_path
.canonicalize()
.map_err(|e| format!("Cannot resolve module path {}: {}", module_path, e))?;
if let Some(cached_module_path) = self.canonical_cache.get(&canonical_path)
&& let Some(nodes) = self.compiled_modules.get(cached_module_path)
{
return Ok(nodes.clone());
}
if self.being_imported.contains(module_path) {
return Err(format!(
"Circular import detected: {} -> {}",
self.import_stack.join(" -> "),
module_path
));
}
self.being_imported.insert(module_path.to_string());
self.import_stack.push(module_path.to_string());
let nodes = self.parse_module(&canonical_path, files, None)?;
self.canonical_cache
.insert(canonical_path, module_path.to_string());
Ok(nodes)
}
pub fn finish_import(&mut self, module_path: &str) {
let cache_key = self.normalize_module_key(module_path).to_string();
self.import_stack.pop();
self.being_imported.remove(&cache_key);
}
pub fn cache_module(&mut self, module_path: &str, nodes: Vec<AstNode>) {
let cache_key = self.normalize_module_key(module_path).to_string();
self.compiled_modules.insert(cache_key, nodes);
}
pub fn check_module_path(&self, module_path: &str) -> (bool, bool) {
if self.has_embedded_module(module_path) {
return (true, false);
}
let embedded_prefix = format!("{}.", module_path);
let has_embedded_directory = self
.embedded_sources
.keys()
.any(|key| key.starts_with(&embedded_prefix));
let mut file_path = self.base_path.clone();
for part in module_path.split('.') {
file_path.push(part);
}
let mux_file = file_path.with_extension("mux");
let dir_path = file_path;
let has_file = mux_file.exists() && mux_file.is_file();
let has_directory = dir_path.exists() && dir_path.is_dir();
(has_file, has_directory || has_embedded_directory)
}
pub fn get_submodules(&self, module_path: &str) -> Result<Vec<String>, String> {
let mut submodules = HashSet::new();
let embedded_prefix = format!("{}.", module_path);
for key in self.embedded_sources.keys() {
if let Some(rest) = key.strip_prefix(&embedded_prefix)
&& !rest.is_empty()
&& let Some(child) = rest.split('.').next()
{
submodules.insert(child.to_string());
}
}
let mut dir_path = self.base_path.clone();
for part in module_path.split('.') {
dir_path.push(part);
}
if dir_path.exists()
&& dir_path.is_dir()
&& let Ok(entries) = std::fs::read_dir(&dir_path)
{
for entry in entries.flatten() {
let path = entry.path();
if path.is_file()
&& path.extension().is_some_and(|ext| ext == "mux")
&& let Some(stem) = path.file_stem().and_then(|s| s.to_str())
{
submodules.insert(stem.to_string());
}
}
}
if submodules.is_empty() {
return Err(format!("Module directory not found: {}", module_path));
}
let mut result: Vec<String> = submodules.into_iter().collect();
result.sort();
Ok(result)
}
fn module_path_to_file(&self, module_path: &str) -> Result<PathBuf, String> {
let mut path = self.base_path.clone();
for part in module_path.split('.') {
path.push(part);
}
path.set_extension("mux");
if !path.exists() {
return Err(format!(
"Module not found: {} (looked for {:?})",
module_path, path
));
}
Ok(path)
}
fn parse_module(
&self,
file_path: &Path,
files: &mut Files,
source_override: Option<&str>,
) -> Result<Vec<AstNode>, String> {
let source_str = if let Some(src) = source_override {
src.to_string()
} else {
std::fs::read_to_string(file_path)
.map_err(|e| format!("Failed to open module: {}", e))?
};
let file_id = files.add(file_path, source_str.clone());
let mut src = Source::from_string(source_str);
let mut lex = Lexer::new(&mut src);
let tokens = match lex.lex_all() {
Ok(t) => t,
Err(e) => {
let emitter = StandardEmitter::new(ColorConfig::Auto);
emitter.emit(&e.to_diagnostic(file_id), files);
return Err(format!("Lexer error in module {}", file_path.display()));
}
};
let mut parser = Parser::new(&tokens);
match parser.parse() {
Ok(nodes) => Ok(nodes),
Err((_, errors)) => {
let emitter = StandardEmitter::new(ColorConfig::Auto);
let diagnostics: Vec<_> = errors.iter().map(|e| e.to_diagnostic(file_id)).collect();
emitter.emit_batch(&diagnostics, files);
Err(format!("Parse error in module {}", file_path.display()))
}
}
}
}